CAGE_expression_inference-apvit/models/AffectNet8_Maxvit_VA/train.py

# import pandas as pd
# import os
# import torch
# import torchvision.transforms as transforms
# import torchvision.models as models
# from torch.utils.data import DataLoader, Dataset
# import torch.nn as nn
# import torch.optim as optim
# from PIL import Image
# from torch.optim import lr_scheduler
# from tqdm import tqdm
# from pathlib import Path

# SCRIPT_DIR = Path(__file__).resolve().parent
# MODEL_PATH = os.getenv("MODEL_PATH", str(SCRIPT_DIR / "model.pt"))


# 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 training and validation from separate CSV files
# 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)
# train_annotations_path = os.getenv(
#     "AFFECTNET_TRAIN_ANNO", f"{ANNO_ROOT}/train_set_annotation_without_lnd.csv"
# )
# valid_annotations_path = os.getenv(
#     "AFFECTNET_VAL_ANNO", f"{ANNO_ROOT}/val_set_annotation_without_lnd.csv"
# )
# train_annotations_df = pd.read_csv(train_annotations_path)
# valid_annotations_df = pd.read_csv(valid_annotations_path)


# # Set parameters
# BATCHSIZE = int(os.getenv("BATCHSIZE", "64"))
# NUM_EPOCHS = 20
# LR = 4e-5
# 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)
# print("MODEL_PATH =", MODEL_PATH)


# # **** 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 = transforms.Compose(
#     [
#         transforms.RandomHorizontalFlip(0.5),
#         transforms.RandomGrayscale(0.01),
#         transforms.RandomRotation(10),
#         transforms.ColorJitter(
#             brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1
#         ),  # model more robust to changes in lighting conditions.
#         transforms.RandomPerspective(
#             distortion_scale=0.2, p=0.5
#         ),  # can be helpful if your images might have varying perspectives.
#         transforms.ToTensor(),  # saves image as tensor (automatically divides by 255)
#         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
#         transforms.RandomErasing(
#             p=0.5, scale=(0.02, 0.2), ratio=(0.3, 3.3), value="random"
#         ),  # Should help overfitting
#     ]
# )

# transform_valid = transforms.Compose(
#     [
#         transforms.ToTensor(),
#         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
#     ]
# )

# train_dataset = CustomDataset(
#     dataframe=train_annotations_df,
#     root_dir=IMAGE_FOLDER,
#     transform=transform,
#     balance=True,
# )
# valid_dataset = CustomDataset(
#     dataframe=valid_annotations_df,
#     root_dir=IMAGE_FOLDER_TEST,
#     transform=transform_valid,
#     balance=False,
# )
# train_loader = DataLoader(
#     train_dataset, batch_size=BATCHSIZE, shuffle=True, num_workers=NUM_WORKERS
# )
# valid_loader = DataLoader(
#     valid_dataset, batch_size=BATCHSIZE, shuffle=False, num_workers=NUM_WORKERS
# )

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

# # Initialize the model
# MODEL = models.maxvit_t(weights="DEFAULT")
# block_channels = MODEL.classifier[3].in_features
# MODEL.classifier = nn.Sequential(
#     nn.AdaptiveAvgPool2d(1),
#     nn.Flatten(),
#     nn.LayerNorm(block_channels),
#     nn.Linear(block_channels, block_channels),
#     nn.Tanh(),
#     nn.Linear(block_channels, 10, bias=False),
# )
# MODEL.to(DEVICE)
# COMBINED_MODEL_PATH = os.getenv(
#     "COMBINED_MODEL_PATH",
#     str((SCRIPT_DIR / "../AffectNet8_Maxvit_Combined/model.pt").resolve()),
# )
# MODEL.load_state_dict(torch.load(COMBINED_MODEL_PATH, map_location=DEVICE))
# MODEL.classifier = nn.Sequential(
#     nn.AdaptiveAvgPool2d(1),
#     nn.Flatten(),
#     nn.LayerNorm(block_channels),
#     nn.Linear(block_channels, block_channels),
#     nn.Tanh(),
#     nn.Dropout(0.3),
#     nn.Linear(block_channels, 2, bias=False),
# )
# MODEL.to(DEVICE)


# def CCCLoss(x, y):
#     # Compute means
#     x_mean = torch.mean(x, dim=0)
#     y_mean = torch.mean(y, dim=0)
#     # Compute variances
#     x_var = torch.var(x, dim=0)
#     y_var = torch.var(y, dim=0)
#     # Compute covariance matrix
#     cov_matrix = torch.matmul(
#         (x - x_mean).permute(*torch.arange(x.dim() - 1, -1, -1)), y - y_mean
#     ) / (x.size(0) - 1)
#     # Compute CCC
#     numerator = 2 * cov_matrix
#     denominator = x_var + y_var + torch.pow((x_mean - y_mean), 2)
#     ccc = torch.mean(numerator / denominator)
#     return -ccc


# val_loss = nn.MSELoss()
# aro_loss = nn.MSELoss()

# optimizer = optim.AdamW(MODEL.parameters(), lr=LR)
# lr_scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=BATCHSIZE * NUM_EPOCHS)

# # ***** Train the model *****
# print("--- Start training ---")
# scaler = torch.cuda.amp.GradScaler()
# best_valid_loss = float("inf")
# l2_lambda = 0.00001  # L1 Regularization
# l1_lambda = 0.00001  # L2 Regularization

# for epoch in range(NUM_EPOCHS):
#     MODEL.train()
#     total_train_correct = 0
#     total_train_samples = 0
#     current_lr = optimizer.param_groups[0]["lr"]
#     for images, _, val_true, aro_true in tqdm(
#         train_loader, desc="Epoch train_loader progress"
#     ):
#         images, val_true, aro_true = (
#             images.to(DEVICE),
#             val_true.to(DEVICE),
#             aro_true.to(DEVICE),
#         )
#         optimizer.zero_grad()
#         train_loss = 0
#         l2_reg = 0
#         l1_reg = 0
#         with torch.autocast(device_type="cuda", dtype=torch.float16):
#             outputs = MODEL(images)
#             val_pred = outputs[:, 0]
#             aro_pred = outputs[:, 1]
#             for param in MODEL.parameters():
#                 l2_reg += torch.norm(param, 2)  # **2
#                 l1_reg += torch.norm(param, 1)
#             loss = (
#                 3 * val_loss(val_pred.cuda(), val_true.cuda())
#                 + 3 * aro_loss(aro_pred.cuda(), aro_true.cuda())
#                 + CCCLoss(val_pred.cuda(), val_true.cuda())
#                 + CCCLoss(aro_pred.cuda(), aro_true.cuda())
#             )
#             # + l2_lambda * l2_reg + l1_lambda * l1_reg
#             train_loss += loss.item()
#             scaler.scale(loss).backward()
#             scaler.step(optimizer)
#             scaler.update()
#             lr_scheduler.step()

#     MODEL.eval()
#     valid_loss = 0.0
#     total_valid_correct = 0
#     total_valid_samples = 0
#     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),
#             )
#             with torch.autocast(device_type="cuda", dtype=torch.float16):
#                 outputs = MODEL(images)
#                 val_pred = outputs[:, 0]
#                 aro_pred = outputs[:, 1]
#                 loss = (
#                     3 * val_loss(val_pred.cuda(), val_true.cuda())
#                     + 3 * aro_loss(aro_pred.cuda(), aro_true.cuda())
#                     + CCCLoss(val_pred.cuda(), val_true.cuda())
#                     + CCCLoss(aro_pred.cuda(), aro_true.cuda())
#                 )
#                 valid_loss += loss.item()

#     print(
#         f"Epoch [{epoch+1}/{NUM_EPOCHS}] - "
#         f"Training Loss: {train_loss/len(train_loader):.4f}, "
#         f"Validation Loss: {valid_loss/len(valid_loader):.4f}, "
#         f"Learning Rate: {current_lr:.8f}, "
#     )

#     valid_loss_mean = valid_loss / max(1, len(valid_loader))

#     if valid_loss_mean < best_valid_loss:
#         best_valid_loss = valid_loss_mean
#         print(f"Saving model at epoch {epoch+1}")
#         torch.save(MODEL.state_dict(), MODEL_PATH)  # Save the best model


# ── 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 numpy as np
import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
import torch.optim as optim
from PIL import Image
from torch.optim import lr_scheduler
from tqdm import tqdm
from pathlib import Path
from sklearn.model_selection import StratifiedKFold
import mmcv
from mmcls.models import build_classifier

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


SCRIPT_DIR = Path(__file__).resolve().parent
MODEL_PATH = os.getenv("MODEL_PATH", str(SCRIPT_DIR / "model.pt"))
PLOTS_DIR = SCRIPT_DIR / "plots"
PLOTS_DIR.mkdir(parents=True, exist_ok=True)
APVIT_REPO = SCRIPT_DIR.parent.parent / "APViT"
APVIT_CONFIG = APVIT_REPO / "configs" / "apvit" / "RAF.py"


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 training and validation from separate CSV files
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)
train_annotations_path = os.getenv(
    "AFFECTNET_TRAIN_ANNO", f"{ANNO_ROOT}/train_set_annotation_without_lnd.csv"
)
valid_annotations_path = os.getenv(
    "AFFECTNET_VAL_ANNO", f"{ANNO_ROOT}/val_set_annotation_without_lnd.csv"
)
train_annotations_df = pd.read_csv(train_annotations_path)
valid_annotations_df = pd.read_csv(valid_annotations_path)


# Set parameters
BATCHSIZE = int(os.getenv("BATCHSIZE", "64"))
NUM_EPOCHS = 20
LR = 4e-5
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)
print("MODEL_PATH =", MODEL_PATH)
print("PLOTS_DIR =", PLOTS_DIR)
print("Train set size =", len(train_annotations_df))
print("Official val_set size (final test only) =", len(valid_annotations_df))
print("StratifiedKFold splits = 5 (each fold train/val = 8:2)")


# =========================
# Plot helpers
# =========================
def save_line_plot(x, ys, labels, title, xlabel, ylabel, save_path):
    plt.figure(figsize=(10, 6))
    for y, label in zip(ys, labels):
        plt.plot(x, y, marker="o", linewidth=2, markersize=4, label=label)
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.savefig(save_path, dpi=200)
    plt.close()


def save_bar_plot(labels, values, title, xlabel, ylabel, save_path, annotate=True):
    plt.figure(figsize=(10, 6))
    bars = plt.bar(labels, values)
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.grid(True, axis="y", alpha=0.3)
    if annotate:
        for bar, val in zip(bars, values):
            plt.text(
                bar.get_x() + bar.get_width() / 2,
                bar.get_height(),
                f"{val:.4f}",
                ha="center",
                va="bottom",
                fontsize=9,
            )
    plt.tight_layout()
    plt.savefig(save_path, dpi=200)
    plt.close()


def save_boxplot(data_list, labels, title, ylabel, save_path):
    plt.figure(figsize=(10, 6))
    plt.boxplot(data_list, labels=labels, vert=True)
    plt.title(title)
    plt.ylabel(ylabel)
    plt.grid(True, axis="y", alpha=0.3)
    plt.tight_layout()
    plt.savefig(save_path, dpi=200)
    plt.close()


def save_histogram(values, bins, title, xlabel, ylabel, save_path):
    plt.figure(figsize=(10, 6))
    plt.hist(values, bins=bins)
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(save_path, dpi=200)
    plt.close()


def save_scatter(x, y, title, xlabel, ylabel, save_path, annotate_labels=None):
    plt.figure(figsize=(10, 6))
    plt.scatter(x, y)
    if annotate_labels is not None:
        for xi, yi, txt in zip(x, y, annotate_labels):
            plt.annotate(str(txt), (xi, yi), fontsize=9)
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(save_path, dpi=200)
    plt.close()


def save_fold_plots(history_df, fold_idx, seed, fold_dir):
    epochs = history_df["epoch"].tolist()

    save_line_plot(
        epochs,
        [history_df["train_loss"].tolist()],
        ["Train Loss"],
        f"Fold {fold_idx} Seed {seed} - Train Loss",
        "Epoch",
        "Loss",
        fold_dir / f"fold_{fold_idx}_train_loss.png",
    )

    save_line_plot(
        epochs,
        [history_df["valid_loss"].tolist()],
        ["Validation Loss"],
        f"Fold {fold_idx} Seed {seed} - Validation Loss",
        "Epoch",
        "Loss",
        fold_dir / f"fold_{fold_idx}_valid_loss.png",
    )

    save_line_plot(
        epochs,
        [history_df["train_loss"].tolist(), history_df["valid_loss"].tolist()],
        ["Train Loss", "Validation Loss"],
        f"Fold {fold_idx} Seed {seed} - Train vs Validation Loss",
        "Epoch",
        "Loss",
        fold_dir / f"fold_{fold_idx}_train_valid_loss.png",
    )

    save_line_plot(
        epochs,
        [history_df["ccc_v"].tolist()],
        ["CCC_v"],
        f"Fold {fold_idx} Seed {seed} - CCC_v",
        "Epoch",
        "CCC",
        fold_dir / f"fold_{fold_idx}_ccc_v.png",
    )

    save_line_plot(
        epochs,
        [history_df["ccc_a"].tolist()],
        ["CCC_a"],
        f"Fold {fold_idx} Seed {seed} - CCC_a",
        "Epoch",
        "CCC",
        fold_dir / f"fold_{fold_idx}_ccc_a.png",
    )

    save_line_plot(
        epochs,
        [history_df["ccc_mean"].tolist()],
        ["CCC_mean"],
        f"Fold {fold_idx} Seed {seed} - CCC_mean",
        "Epoch",
        "CCC",
        fold_dir / f"fold_{fold_idx}_ccc_mean.png",
    )

    save_line_plot(
        epochs,
        [
            history_df["ccc_v"].tolist(),
            history_df["ccc_a"].tolist(),
            history_df["ccc_mean"].tolist(),
        ],
        ["CCC_v", "CCC_a", "CCC_mean"],
        f"Fold {fold_idx} Seed {seed} - Combined CCC Curves",
        "Epoch",
        "CCC",
        fold_dir / f"fold_{fold_idx}_ccc_all.png",
    )

    save_line_plot(
        epochs,
        [history_df["lr"].tolist()],
        ["Learning Rate"],
        f"Fold {fold_idx} Seed {seed} - Learning Rate",
        "Epoch",
        "LR",
        fold_dir / f"fold_{fold_idx}_lr.png",
    )

    save_scatter(
        history_df["train_loss"].tolist(),
        history_df["valid_loss"].tolist(),
        f"Fold {fold_idx} Seed {seed} - Train Loss vs Valid Loss",
        "Train Loss",
        "Validation Loss",
        fold_dir / f"fold_{fold_idx}_train_vs_valid_scatter.png",
        annotate_labels=epochs,
    )

    save_scatter(
        history_df["ccc_v"].tolist(),
        history_df["ccc_a"].tolist(),
        f"Fold {fold_idx} Seed {seed} - CCC_v vs CCC_a",
        "CCC_v",
        "CCC_a",
        fold_dir / f"fold_{fold_idx}_ccc_v_vs_ccc_a.png",
        annotate_labels=epochs,
    )


def save_summary_plots(summary_df, save_dir):
    fold_labels = [f"Fold {x}" for x in summary_df["fold"].tolist()]

    save_bar_plot(
        fold_labels,
        summary_df["test_loss"].tolist(),
        "Final Test Loss by Fold",
        "Fold",
        "Loss",
        save_dir / "summary_test_loss_by_fold.png",
    )

    save_bar_plot(
        fold_labels,
        summary_df["test_ccc_v"].tolist(),
        "Final Test CCC_v by Fold",
        "Fold",
        "CCC_v",
        save_dir / "summary_test_ccc_v_by_fold.png",
    )

    save_bar_plot(
        fold_labels,
        summary_df["test_ccc_a"].tolist(),
        "Final Test CCC_a by Fold",
        "Fold",
        "CCC_a",
        save_dir / "summary_test_ccc_a_by_fold.png",
    )

    save_bar_plot(
        fold_labels,
        summary_df["test_ccc_mean"].tolist(),
        "Final Test CCC_mean by Fold",
        "Fold",
        "CCC_mean",
        save_dir / "summary_test_ccc_mean_by_fold.png",
    )

    save_bar_plot(
        fold_labels,
        summary_df["best_epoch"].tolist(),
        "Best Epoch by Fold",
        "Fold",
        "Best Epoch",
        save_dir / "summary_best_epoch_by_fold.png",
        annotate=True,
    )

    save_boxplot(
        [
            summary_df["test_ccc_v"].tolist(),
            summary_df["test_ccc_a"].tolist(),
            summary_df["test_ccc_mean"].tolist(),
        ],
        ["CCC_v", "CCC_a", "CCC_mean"],
        "Distribution of Final CCC Across 5 Folds",
        "CCC",
        save_dir / "summary_ccc_boxplot.png",
    )

    save_histogram(
        summary_df["test_ccc_mean"].tolist(),
        bins=min(5, len(summary_df)),
        title="Histogram of Final CCC_mean",
        xlabel="CCC_mean",
        ylabel="Frequency",
        save_path=save_dir / "summary_ccc_mean_histogram.png",
    )

    save_scatter(
        summary_df["best_epoch"].tolist(),
        summary_df["test_ccc_mean"].tolist(),
        "Best Epoch vs Final CCC_mean",
        "Best Epoch",
        "Final CCC_mean",
        save_dir / "summary_best_epoch_vs_ccc_mean.png",
        annotate_labels=summary_df["fold"].tolist(),
    )


# **** 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):
        img_id = int(float(self.dataframe["number"].iloc[idx]))
        image_path = os.path.join(self.root_dir, f"{img_id}.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 = transforms.Compose(
    [
        transforms.Resize(112),  # APViT / IR-50 requires 112x112 input
        transforms.RandomHorizontalFlip(0.5),
        transforms.RandomGrayscale(0.01),
        transforms.RandomRotation(10),
        transforms.ColorJitter(
            brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1
        ),
        transforms.RandomPerspective(
            distortion_scale=0.2, p=0.5
        ),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        transforms.RandomErasing(
            p=0.5, scale=(0.02, 0.2), ratio=(0.3, 3.3), value="random"
        ),
    ]
)

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]),
    ]
)

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


APVIT_EMBED_DIM = 768


def _build_apvit_backbone():
    """Build APViT PoolingVitClassifier (backbone only, no head)."""
    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
    if cfg.model.extractor.pretrained is None:
        print("[WARNING] IR-50 weights not found at", ir50_w, "— training from random init.")
    if cfg.model.vit.pretrained is None:
        print("[WARNING] ViT-Small weights not found at", vit_small_w, "— training from random init.")
    return build_classifier(cfg.model)


class APViTWithHead(nn.Module):
    def __init__(self, num_outputs):
        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.Linear(APVIT_EMBED_DIM, num_outputs, bias=False),
        )

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


COMBINED_MODEL_PATH = os.getenv(
    "COMBINED_MODEL_PATH",
    str((SCRIPT_DIR / "../AffectNet8_Maxvit_Combined/model.pt").resolve()),
)


def build_model():
    model = APViTWithHead(num_outputs=10)  # Match Combined checkpoint structure
    model.to(DEVICE)
    model.load_state_dict(torch.load(COMBINED_MODEL_PATH, map_location=DEVICE))
    # Replace head with VA regression head (2 outputs)
    model.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),
    )
    model.to(DEVICE)
    return model


def CCCLoss(x, y):
    x_mean = torch.mean(x, dim=0)
    y_mean = torch.mean(y, dim=0)
    x_var = torch.var(x, dim=0)
    y_var = torch.var(y, dim=0)
    cov_matrix = torch.matmul(
        (x - x_mean).permute(*torch.arange(x.dim() - 1, -1, -1)), y - y_mean
    ) / (x.size(0) - 1)
    numerator = 2 * cov_matrix
    denominator = x_var + y_var + torch.pow((x_mean - y_mean), 2)
    ccc = torch.mean(numerator / denominator)
    return -ccc


def ccc_numpy_1d(x, y, eps=1e-8):
    x = np.asarray(x).reshape(-1)
    y = np.asarray(y).reshape(-1)
    x_mean = x.mean()
    y_mean = y.mean()
    x_var = x.var()
    y_var = y.var()
    cov = ((x - x_mean) * (y - y_mean)).mean()
    return (2.0 * cov) / (x_var + y_var + (x_mean - y_mean) ** 2 + eps)


def evaluate_loader(model, loader):
    model.eval()
    eval_loss = 0.0
    eval_val_true = []
    eval_val_pred = []
    eval_aro_true = []
    eval_aro_pred = []
    with torch.no_grad():
        for images, _, val_true, aro_true in loader:
            images, val_true, aro_true = (
                images.to(DEVICE),
                val_true.to(DEVICE),
                aro_true.to(DEVICE),
            )
            with torch.autocast(device_type="cuda", dtype=torch.float16):
                outputs = model(images)
                val_pred = outputs[:, 0]
                aro_pred = outputs[:, 1]
                loss = (
                    3 * val_loss(val_pred.cuda(), val_true.cuda())
                    + 3 * aro_loss(aro_pred.cuda(), aro_true.cuda())
                    + CCCLoss(val_pred.cuda(), val_true.cuda())
                    + CCCLoss(aro_pred.cuda(), aro_true.cuda())
                )
                eval_loss += loss.item()

            eval_val_true.extend(val_true.detach().cpu().float().numpy())
            eval_val_pred.extend(val_pred.detach().cpu().float().numpy())
            eval_aro_true.extend(aro_true.detach().cpu().float().numpy())
            eval_aro_pred.extend(aro_pred.detach().cpu().float().numpy())

    ccc_v = ccc_numpy_1d(eval_val_pred, eval_val_true)
    ccc_a = ccc_numpy_1d(eval_aro_pred, eval_aro_true)
    ccc_mean = (ccc_v + ccc_a) / 2.0
    return eval_loss / max(1, len(loader)), ccc_v, ccc_a, ccc_mean


val_loss = nn.MSELoss()
aro_loss = nn.MSELoss()

# ***** Train the model *****
print("--- Start training on full train_set, validate on val_set ---")
l2_lambda = 0.00001
l1_lambda = 0.00001

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

model = build_model()
optimizer = optim.AdamW(model.parameters(), lr=LR)
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=BATCHSIZE * NUM_EPOCHS)
scaler = torch.cuda.amp.GradScaler()
best_valid_loss = float("inf")
best_epoch = -1
history = {
    "epoch": [],
    "train_loss": [],
    "valid_loss": [],
    "ccc_v": [],
    "ccc_a": [],
    "ccc_mean": [],
    "lr": [],
}

for epoch in range(NUM_EPOCHS):
    model.train()
    train_loss = 0.0
    train_val_true = []
    train_val_pred = []
    train_aro_true = []
    train_aro_pred = []
    current_lr = optimizer.param_groups[0]["lr"]
    for images, _, val_true, aro_true in tqdm(train_loader, desc=f"Epoch {epoch+1} train_loader progress"):
        images, val_true, aro_true = images.to(DEVICE), val_true.to(DEVICE), aro_true.to(DEVICE)
        optimizer.zero_grad()
        l2_reg = 0
        l1_reg = 0
        with torch.autocast(device_type="cuda", dtype=torch.float16):
            outputs = model(images)
            val_pred = outputs[:, 0]
            aro_pred = outputs[:, 1]
            for param in model.parameters():
                l2_reg += torch.norm(param, 2)
                l1_reg += torch.norm(param, 1)
            loss = (
                3 * val_loss(val_pred, val_true)
                + 3 * aro_loss(aro_pred, aro_true)
                + CCCLoss(val_pred, val_true)
                + CCCLoss(aro_pred, aro_true)
            )
            train_loss += loss.item()
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()
            train_val_true.extend(val_true.detach().cpu().numpy())
            train_val_pred.extend(val_pred.detach().cpu().numpy())
            train_aro_true.extend(aro_true.detach().cpu().numpy())
            train_aro_pred.extend(aro_pred.detach().cpu().numpy())

    train_loss /= max(1, len(train_loader))
    train_ccc_v = ccc_numpy_1d(train_val_pred, train_val_true)
    train_ccc_a = ccc_numpy_1d(train_aro_pred, train_aro_true)
    train_ccc_mean = (train_ccc_v + train_ccc_a) / 2.0

    model.eval()
    valid_loss = 0.0
    valid_val_true = []
    valid_val_pred = []
    valid_aro_true = []
    valid_aro_pred = []
    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)
            with torch.autocast(device_type="cuda", dtype=torch.float16):
                outputs = model(images)
                val_pred = outputs[:, 0]
                aro_pred = outputs[:, 1]
                loss = (
                    3 * val_loss(val_pred, val_true)
                    + 3 * aro_loss(aro_pred, aro_true)
                    + CCCLoss(val_pred, val_true)
                    + CCCLoss(aro_pred, aro_true)
                )
                valid_loss += loss.item()
                valid_val_true.extend(val_true.detach().cpu().numpy())
                valid_val_pred.extend(val_pred.detach().cpu().numpy())
                valid_aro_true.extend(aro_true.detach().cpu().numpy())
                valid_aro_pred.extend(aro_pred.detach().cpu().numpy())

    valid_loss /= max(1, len(valid_loader))
    valid_ccc_v = ccc_numpy_1d(valid_val_pred, valid_val_true)
    valid_ccc_a = ccc_numpy_1d(valid_aro_pred, valid_aro_true)
    valid_ccc_mean = (valid_ccc_v + valid_ccc_a) / 2.0

    history["epoch"].append(epoch + 1)
    history["train_loss"].append(train_loss)
    history["valid_loss"].append(valid_loss)
    history["ccc_v"].append(valid_ccc_v)
    history["ccc_a"].append(valid_ccc_a)
    history["ccc_mean"].append(valid_ccc_mean)
    history["lr"].append(current_lr)

    print(
        f"Epoch [{epoch+1}/{NUM_EPOCHS}] - "
        f"Train Loss: {train_loss:.4f}, "
        f"Valid Loss: {valid_loss:.4f}, "
        f"Valid CCC_v: {valid_ccc_v:.4f}, "
        f"Valid CCC_a: {valid_ccc_a:.4f}, "
        f"Valid CCC_mean: {valid_ccc_mean:.4f}, "
        f"Learning Rate: {current_lr:.8f}"
    )

    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        best_epoch = epoch + 1
        print(f"Saving best model at epoch {best_epoch}")
        torch.save(model.state_dict(), MODEL_PATH)

# Save training history and plot if needed
import pandas as pd
history_df = pd.DataFrame(history)
history_df.to_csv(PLOTS_DIR / "train_history.csv", index=False)

# Plotting (reuse your plot helpers, but only plot what makes sense)
save_line_plot(
    history_df["epoch"].tolist(),
    [history_df["train_loss"].tolist()],
    ["Train Loss"],
    "Train Loss",
    "Epoch",
    "Loss",
    PLOTS_DIR / "train_loss.png",
)
save_line_plot(
    history_df["epoch"].tolist(),
    [history_df["valid_loss"].tolist()],
    ["Validation Loss"],
    "Validation Loss",
    "Epoch",
    "Loss",
    PLOTS_DIR / "valid_loss.png",
)
save_line_plot(
    history_df["epoch"].tolist(),
    [history_df["train_loss"].tolist(), history_df["valid_loss"].tolist()],
    ["Train Loss", "Validation Loss"],
    "Train vs Validation Loss",
    "Epoch",
    "Loss",
    PLOTS_DIR / "train_vs_valid_loss.png",
)
save_line_plot(
    history_df["epoch"].tolist(),
    [history_df["ccc_v"].tolist()],
    ["CCC_v"],
    "Validation CCC_v",
    "Epoch",
    "CCC",
    PLOTS_DIR / "valid_ccc_v.png",
)
save_line_plot(
    history_df["epoch"].tolist(),
    [history_df["ccc_a"].tolist()],
    ["CCC_a"],
    "Validation CCC_a",
    "Epoch",
    "CCC",
    PLOTS_DIR / "valid_ccc_a.png",
)
save_line_plot(
    history_df["epoch"].tolist(),
    [history_df["ccc_mean"].tolist()],
    ["CCC_mean"],
    "Validation CCC_mean",
    "Epoch",
    "CCC",
    PLOTS_DIR / "valid_ccc_mean.png",
)
save_line_plot(
    history_df["epoch"].tolist(),
    [history_df["lr"].tolist()],
    ["Learning Rate"],
    "Learning Rate",
    "Epoch",
    "LR",
    PLOTS_DIR / "lr.png",
)


# In thông tin của epoch tốt nhất
best_idx = best_epoch - 1
print("=== Best Epoch Metrics ===")
print(f"Best epoch: {best_epoch}")
print(f"Train Loss: {history_df['train_loss'][best_idx]:.4f}")
print(f"Valid Loss: {history_df['valid_loss'][best_idx]:.4f}")
print(f"Valid CCC_v: {history_df['ccc_v'][best_idx]:.4f}")
print(f"Valid CCC_a: {history_df['ccc_a'][best_idx]:.4f}")
print(f"Valid CCC_mean: {history_df['ccc_mean'][best_idx]:.4f}")

n_epochs_run = len(history_df["ccc_v"])
avg_ccc_v = sum(history_df["ccc_v"]) / n_epochs_run
avg_ccc_a = sum(history_df["ccc_a"]) / n_epochs_run
avg_ccc_mean = sum(history_df["ccc_mean"]) / n_epochs_run
avg_valid_loss = sum(history_df["valid_loss"]) / n_epochs_run
print("\n=== Average Across All Epochs ===")
print(f"Avg Valid Loss  : {avg_valid_loss:.4f}")
print(f"Avg CCC_v       : {avg_ccc_v:.4f}")
print(f"Avg CCC_a       : {avg_ccc_a:.4f}")
print(f"Avg CCC_mean    : {avg_ccc_mean:.4f}")
print("=================================")

print(f"--- Training complete. Best epoch: {best_epoch}, Best val loss: {best_valid_loss:.4f} ---")
print(f"Best model saved to: {MODEL_PATH}")
print(f"All plots and CSV logs saved to: {PLOTS_DIR}")