huantran2004

Upload 127 files

68efce6 verified 6 months ago

6.82 kB

# ── APViT repo path setup (must precede mmcls/mmcv imports) ─────────────────
import sys
from pathlib import Path as _Path
_apvit_path = str(_Path(__file__).resolve().parent.parent.parent / "APViT")
if _apvit_path not in sys.path:
    sys.path.insert(0, _apvit_path)
# ─────────────────────────────────────────────────────────────────────────────

import pandas as pd
import os
import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
from PIL import Image
from pathlib import Path
import mmcv
from mmcls.models import build_classifier


def resolve_images_dir(split_name, explicit_env_var, data_root, extract_root):
    explicit = os.getenv(explicit_env_var, "").strip()
    candidates = []
    if explicit:
        candidates.append(Path(explicit))

    candidates.append(Path(data_root) / f"{split_name}_set" / "images")
    candidates.append(Path(extract_root) / f"{split_name}_extracted" / "images")

    extract_split_root = Path(extract_root) / f"{split_name}_extracted"
    if extract_split_root.exists():
        for p in extract_split_root.rglob("images"):
            if p.is_dir():
                candidates.append(p)

    for p in candidates:
        if p.exists() and p.is_dir():
            return str(p)

    tried = "\n".join([str(p) for p in candidates])
    raise FileNotFoundError(
        f"Could not find images folder for split='{split_name}'. Tried:\n{tried}"
    )

# Load the annotations for validation from CSV file
DATA_ROOT = os.getenv("AFFECTNET_ROOT", "/workspace/data_affectnet/AffectNet")
EXTRACT_ROOT = os.getenv("AFFECTNET_EXTRACT_ROOT", f"{DATA_ROOT}/extracted")
ANNO_ROOT = os.getenv("AFFECTNET_ANNO_ROOT", "../../affectnet_annotations")
IMAGE_FOLDER = resolve_images_dir("train", "AFFECTNET_TRAIN_IMAGES", DATA_ROOT, EXTRACT_ROOT)
IMAGE_FOLDER_TEST = resolve_images_dir("val", "AFFECTNET_VAL_IMAGES", DATA_ROOT, EXTRACT_ROOT)
valid_annotations_path = os.getenv(
    "AFFECTNET_VAL_ANNO", f"{ANNO_ROOT}/val_set_annotation_without_lnd.csv"
)
valid_annotations_df = pd.read_csv(valid_annotations_path)

# Set parameters
BATCHSIZE = int(os.getenv("BATCHSIZE", "64"))
NUM_WORKERS = int(os.getenv("NUM_WORKERS", "0"))
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("DATA_ROOT =", DATA_ROOT)
print("EXTRACT_ROOT =", EXTRACT_ROOT)
print("AFFECTNET_TRAIN_IMAGES =", IMAGE_FOLDER)
print("AFFECTNET_VAL_IMAGES =", IMAGE_FOLDER_TEST)


# **** Create dataset and data loaders ****
class CustomDataset(Dataset):
    def __init__(self, dataframe, root_dir, transform=None, balance=False):
        self.dataframe = dataframe
        self.transform = transform
        self.root_dir = root_dir
        self.balance = balance

        if self.balance:
            self.dataframe = self.balance_dataset()

    def __len__(self):
        return len(self.dataframe)

    def __getitem__(self, idx):
        image_path = os.path.join(
            self.root_dir, f"{self.dataframe['number'].iloc[idx]}.jpg"
        )
        image = Image.open(image_path)

        classes = torch.tensor(self.dataframe.iloc[idx, 1], dtype=torch.int8)
        valence = torch.tensor(self.dataframe.iloc[idx, 2], dtype=torch.float16)
        arousal = torch.tensor(self.dataframe.iloc[idx, 3], dtype=torch.float16)

        if self.transform:
            image = self.transform(image)

        return image, classes, valence, arousal

    def balance_dataset(self):
        balanced_df = self.dataframe.groupby("exp", group_keys=False).apply(
            lambda x: x.sample(self.dataframe["exp"].value_counts().min())
        )
        return balanced_df


transform_valid = transforms.Compose(
    [
        transforms.Resize(112),  # APViT / IR-50 requires 112x112 input
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)

valid_dataset = CustomDataset(
    dataframe=valid_annotations_df,
    root_dir=IMAGE_FOLDER_TEST,
    transform=transform_valid,
    balance=False,
)
valid_loader = DataLoader(
    valid_dataset, batch_size=BATCHSIZE, shuffle=False, num_workers=NUM_WORKERS
)

# ***** Define the model *****

APVIT_EMBED_DIM = 768


def _build_apvit_backbone():
    cfg = mmcv.Config.fromfile(str(APVIT_CONFIG))
    cfg.model.pretrained = None
    cfg.model.head = None
    ir50_w = APVIT_REPO / "weights" / "backbone_ir50_ms1m_epoch120.pth"
    vit_small_w = APVIT_REPO / "weights" / "vit_small_p16_224-15ec54c9.pth"
    cfg.model.extractor.pretrained = str(ir50_w) if ir50_w.exists() else None
    cfg.model.vit.pretrained = str(vit_small_w) if vit_small_w.exists() else None
    return build_classifier(cfg.model)


class APViTVAModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.apvit = _build_apvit_backbone()
        self.head = nn.Sequential(
            nn.LayerNorm(APVIT_EMBED_DIM),
            nn.Linear(APVIT_EMBED_DIM, APVIT_EMBED_DIM),
            nn.Tanh(),
            nn.Dropout(0.3),
            nn.Linear(APVIT_EMBED_DIM, 2, bias=False),
        )

    def forward(self, x):
        features, _ = self.apvit.extract_feat(x)  # [B, 768] CLS token
        return self.head(features)


# Initialize the model
MODEL = APViTVAModel()
MODEL.to(DEVICE)

# **** Test the model performance for classification ****

# Set the model to evaluation mode
MODEL.load_state_dict(torch.load("model.pt", map_location=DEVICE))
MODEL.to(DEVICE)
MODEL.eval()

all_val_true_values = []
all_val_predicted_values = []
all_aro_true_values = []
all_aro_predicted_values = []

# Start inference on test set
with torch.no_grad():
    for images, _, val_true, aro_true in valid_loader:
        images, val_true, aro_true = (
            images.to(DEVICE),
            val_true.to(DEVICE),
            aro_true.to(DEVICE),
        )

        outputs = MODEL(images)
        val_pred = outputs[:, 0]
        aro_pred = outputs[:, 1]

        # Append to the lists --> Regression
        true_val_values = val_true.cpu().numpy()
        true_aro_values = aro_true.cpu().numpy()
        pred_val_values = val_pred.cpu().numpy()
        pred_aro_values = aro_pred.cpu().numpy()
        all_val_true_values.extend(true_val_values)
        all_aro_true_values.extend(true_aro_values)
        all_val_predicted_values.extend(pred_val_values)
        all_aro_predicted_values.extend(pred_aro_values)
df = pd.DataFrame(
    {
        "val_pred": all_val_predicted_values,
        "val_true": all_val_true_values,
        "aro_pred": all_aro_predicted_values,
        "aro_true": all_aro_true_values,
    }
)
df.to_csv("inference.csv", index=False)

	# ── APViT repo path setup (must precede mmcls/mmcv imports) ─────────────────
	import sys
	from pathlib import Path as _Path
	_apvit_path = str(_Path(__file__).resolve().parent.parent.parent / "APViT")
	if _apvit_path not in sys.path:
	sys.path.insert(0, _apvit_path)
	# ─────────────────────────────────────────────────────────────────────────────

	import pandas as pd
	import os
	import torch
	import torchvision.transforms as transforms
	from torch.utils.data import DataLoader, Dataset
	import torch.nn as nn
	from PIL import Image
	from pathlib import Path
	import mmcv
	from mmcls.models import build_classifier


	def resolve_images_dir(split_name, explicit_env_var, data_root, extract_root):
	explicit = os.getenv(explicit_env_var, "").strip()
	candidates = []
	if explicit:
	candidates.append(Path(explicit))

	candidates.append(Path(data_root) / f"{split_name}_set" / "images")
	candidates.append(Path(extract_root) / f"{split_name}_extracted" / "images")

	extract_split_root = Path(extract_root) / f"{split_name}_extracted"
	if extract_split_root.exists():
	for p in extract_split_root.rglob("images"):
	if p.is_dir():
	candidates.append(p)

	for p in candidates:
	if p.exists() and p.is_dir():
	return str(p)

	tried = "\n".join([str(p) for p in candidates])
	raise FileNotFoundError(
	f"Could not find images folder for split='{split_name}'. Tried:\n{tried}"
	)

	# Load the annotations for validation from CSV file
	DATA_ROOT = os.getenv("AFFECTNET_ROOT", "/workspace/data_affectnet/AffectNet")
	EXTRACT_ROOT = os.getenv("AFFECTNET_EXTRACT_ROOT", f"{DATA_ROOT}/extracted")
	ANNO_ROOT = os.getenv("AFFECTNET_ANNO_ROOT", "../../affectnet_annotations")
	IMAGE_FOLDER = resolve_images_dir("train", "AFFECTNET_TRAIN_IMAGES", DATA_ROOT, EXTRACT_ROOT)
	IMAGE_FOLDER_TEST = resolve_images_dir("val", "AFFECTNET_VAL_IMAGES", DATA_ROOT, EXTRACT_ROOT)
	valid_annotations_path = os.getenv(
	"AFFECTNET_VAL_ANNO", f"{ANNO_ROOT}/val_set_annotation_without_lnd.csv"
	)
	valid_annotations_df = pd.read_csv(valid_annotations_path)

	# Set parameters
	BATCHSIZE = int(os.getenv("BATCHSIZE", "64"))
	NUM_WORKERS = int(os.getenv("NUM_WORKERS", "0"))
	DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

	print("DATA_ROOT =", DATA_ROOT)
	print("EXTRACT_ROOT =", EXTRACT_ROOT)
	print("AFFECTNET_TRAIN_IMAGES =", IMAGE_FOLDER)
	print("AFFECTNET_VAL_IMAGES =", IMAGE_FOLDER_TEST)


	# ** Create dataset and data loaders **
	class CustomDataset(Dataset):
	def __init__(self, dataframe, root_dir, transform=None, balance=False):
	self.dataframe = dataframe
	self.transform = transform
	self.root_dir = root_dir
	self.balance = balance

	if self.balance:
	self.dataframe = self.balance_dataset()

	def __len__(self):
	return len(self.dataframe)

	def __getitem__(self, idx):
	image_path = os.path.join(
	self.root_dir, f"{self.dataframe['number'].iloc[idx]}.jpg"
	)
	image = Image.open(image_path)

	classes = torch.tensor(self.dataframe.iloc[idx, 1], dtype=torch.int8)
	valence = torch.tensor(self.dataframe.iloc[idx, 2], dtype=torch.float16)
	arousal = torch.tensor(self.dataframe.iloc[idx, 3], dtype=torch.float16)

	if self.transform:
	image = self.transform(image)

	return image, classes, valence, arousal

	def balance_dataset(self):
	balanced_df = self.dataframe.groupby("exp", group_keys=False).apply(
	lambda x: x.sample(self.dataframe["exp"].value_counts().min())
	)
	return balanced_df


	transform_valid = transforms.Compose(
	[
	transforms.Resize(112), # APViT / IR-50 requires 112x112 input
	transforms.ToTensor(),
	transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
	]
	)

	valid_dataset = CustomDataset(
	dataframe=valid_annotations_df,
	root_dir=IMAGE_FOLDER_TEST,
	transform=transform_valid,
	balance=False,
	)
	valid_loader = DataLoader(
	valid_dataset, batch_size=BATCHSIZE, shuffle=False, num_workers=NUM_WORKERS
	)

	# *** Define the model ***

	APVIT_EMBED_DIM = 768


	def _build_apvit_backbone():
	cfg = mmcv.Config.fromfile(str(APVIT_CONFIG))
	cfg.model.pretrained = None
	cfg.model.head = None
	ir50_w = APVIT_REPO / "weights" / "backbone_ir50_ms1m_epoch120.pth"
	vit_small_w = APVIT_REPO / "weights" / "vit_small_p16_224-15ec54c9.pth"
	cfg.model.extractor.pretrained = str(ir50_w) if ir50_w.exists() else None
	cfg.model.vit.pretrained = str(vit_small_w) if vit_small_w.exists() else None
	return build_classifier(cfg.model)


	class APViTVAModel(nn.Module):
	def __init__(self):
	super().__init__()
	self.apvit = _build_apvit_backbone()
	self.head = nn.Sequential(
	nn.LayerNorm(APVIT_EMBED_DIM),
	nn.Linear(APVIT_EMBED_DIM, APVIT_EMBED_DIM),
	nn.Tanh(),
	nn.Dropout(0.3),
	nn.Linear(APVIT_EMBED_DIM, 2, bias=False),
	)

	def forward(self, x):
	features, _ = self.apvit.extract_feat(x) # [B, 768] CLS token
	return self.head(features)


	# Initialize the model
	MODEL = APViTVAModel()
	MODEL.to(DEVICE)

	# ** Test the model performance for classification **

	# Set the model to evaluation mode
	MODEL.load_state_dict(torch.load("model.pt", map_location=DEVICE))
	MODEL.to(DEVICE)
	MODEL.eval()

	all_val_true_values = []
	all_val_predicted_values = []
	all_aro_true_values = []
	all_aro_predicted_values = []

	# Start inference on test set
	with torch.no_grad():
	for images, _, val_true, aro_true in valid_loader:
	images, val_true, aro_true = (
	images.to(DEVICE),
	val_true.to(DEVICE),
	aro_true.to(DEVICE),
	)

	outputs = MODEL(images)
	val_pred = outputs[:, 0]
	aro_pred = outputs[:, 1]

	# Append to the lists --> Regression
	true_val_values = val_true.cpu().numpy()
	true_aro_values = aro_true.cpu().numpy()
	pred_val_values = val_pred.cpu().numpy()
	pred_aro_values = aro_pred.cpu().numpy()
	all_val_true_values.extend(true_val_values)
	all_aro_true_values.extend(true_aro_values)
	all_val_predicted_values.extend(pred_val_values)
	all_aro_predicted_values.extend(pred_aro_values)
	df = pd.DataFrame(
	{
	"val_pred": all_val_predicted_values,
	"val_true": all_val_true_values,
	"aro_pred": all_aro_predicted_values,
	"aro_true": all_aro_true_values,
	}
	)
	df.to_csv("inference.csv", index=False)