{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"accelerator":"GPU","colab":{"collapsed_sections":[],"name":"classification.ipynb","toc_visible":true},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":22881,"databundleVersionId":1552852,"sourceType":"competition"}],"dockerImageVersionId":31153,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"new arcface\n","metadata":{}},{"cell_type":"code","source":"import os\nimport json\nimport math\nimport random\nimport numpy as np\nfrom PIL import Image\nfrom collections import defaultdict\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader, Sampler\nfrom torch.cuda.amp import autocast, GradScaler\n\nimport torchvision.transforms as T\nimport torchvision.models as models\n\nfrom sklearn.metrics import roc_curve, auc\n\n# ============================================================\n# CONFIG\n# ============================================================\n\nIMG_SIZE = 160\nEMBEDDING_DIM = 512\nBATCH_IDENTITIES = 16\nIMAGES_PER_ID = 4\nBATCH_SIZE = BATCH_IDENTITIES * IMAGES_PER_ID\n\nEPOCHS_STAGE1 = 40\nEPOCHS_STAGE2 = 20\nWARMUP_EPOCHS = 2\n\nBASE_LR = 3e-4\nWEIGHT_DECAY = 1e-4\n\nARCFACE_MARGIN = 0.5\nARCFACE_SCALE = 64.0\nTRIPLET_WEIGHT = 0.1\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n# ============================================================\n# PATHS\n# ============================================================\n\nBASE_DIR = '/kaggle/input/11-785-fall-20-homework-2-part-2'\nTRAIN_DATASET_PATH = os.path.join(BASE_DIR, \"classification_data/train_data\")\nVAL_DATASET_PATH = os.path.join(BASE_DIR, \"classification_data/val_data\")\nVAL_PAIRS_FILE = os.path.join(BASE_DIR, \"verification_pairs_val.txt\")\nOUTPUT_DIR = \"/kaggle/working/\"\nOUTPUT_MODEL_PATH = os.path.join(OUTPUT_DIR, \"best_model.pth\")\nOUTPUT_METADATA_PATH = os.path.join(OUTPUT_DIR, \"metadata.json\")\n\n# ============================================================\n# DATASETS\n# ============================================================\n\nclass FaceFolderDataset(Dataset):\n    def __init__(self, root, transform=None):\n        self.samples = []\n        self.labels = {}\n        self.transform = transform\n\n        for idx, person in enumerate(sorted(os.listdir(root))):\n            person_dir = os.path.join(root, person)\n            if not os.path.isdir(person_dir):\n                continue\n            self.labels[person] = idx\n            for img in os.listdir(person_dir):\n                self.samples.append(\n                    (os.path.join(person_dir, img), idx)\n                )\n\n    def __len__(self):\n        return len(self.samples)\n\n    def __getitem__(self, idx):\n        path, label = self.samples[idx]\n        img = Image.open(path).convert(\"RGB\")\n        if self.transform:\n            img = self.transform(img)\n        return img, label\n\n\nclass BalancedIdentitySampler(Sampler):\n    def __init__(self, labels, n_ids, n_imgs):\n        self.labels = labels\n        self.n_ids = n_ids\n        self.n_imgs = n_imgs\n\n        self.index_dict = defaultdict(list)\n        for idx, label in enumerate(labels):\n            self.index_dict[label].append(idx)\n\n        self.labels_set = list(self.index_dict.keys())\n\n    def __iter__(self):\n        random.shuffle(self.labels_set)\n        batch = []\n\n        for label in self.labels_set:\n            imgs = random.sample(\n                self.index_dict[label],\n                min(len(self.index_dict[label]), self.n_imgs)\n            )\n            batch.extend(imgs)\n            if len(batch) == self.n_ids * self.n_imgs:\n                yield from batch\n                batch = []\n\n    def __len__(self):\n        return len(self.labels)\n\n\n# ============================================================\n# MODEL\n# ============================================================\n\nclass FaceNet(nn.Module):\n    def __init__(self, embedding_dim):\n        super().__init__()\n        backbone = models.mobilenet_v3_large(pretrained=True)\n        backbone.classifier = nn.Identity()\n        self.backbone = backbone\n        self.fc = nn.Linear(960, embedding_dim)\n\n    def forward(self, x):\n        x = self.backbone(x)\n        x = self.fc(x)\n        return F.normalize(x)\n\n\nclass ArcFace(nn.Module):\n    def __init__(self, embedding_dim, num_classes, margin, scale):\n        super().__init__()\n        self.weight = nn.Parameter(torch.randn(num_classes, embedding_dim))\n        nn.init.xavier_uniform_(self.weight)\n        self.margin = margin\n        self.scale = scale\n\n    def forward(self, embeddings, labels):\n        cosine = F.linear(embeddings, F.normalize(self.weight))\n        theta = torch.acos(torch.clamp(cosine, -1 + 1e-7, 1 - 1e-7))\n        target_logits = torch.cos(theta + self.margin)\n\n        one_hot = torch.zeros_like(cosine)\n        one_hot.scatter_(1, labels.view(-1, 1), 1.0)\n\n        output = cosine * (1 - one_hot) + target_logits * one_hot\n        return output * self.scale\n\n\n# ============================================================\n# LOSSES\n# ============================================================\n\ndef triplet_loss(embeddings, labels, margin=0.3):\n    dist = torch.cdist(embeddings, embeddings)\n    loss = 0.0\n    count = 0\n\n    for i in range(len(embeddings)):\n        pos = dist[i][labels == labels[i]]\n        neg = dist[i][labels != labels[i]]\n        if len(pos) > 1 and len(neg) > 0:\n            loss += F.relu(pos.max() - neg.min() + margin)\n            count += 1\n\n    return loss / max(count, 1)\n\n\n# ============================================================\n# VALIDATION (ROC / EER)\n# ============================================================\n\n# def evaluate(model, pairs_file, root, transform):\n#     model.eval()\n#     sims, labels = [], []\n\n#     with torch.no_grad():\n#         for line in open(pairs_file):\n#             p1, p2, lab = line.strip().split()\n#             img1 = transform(Image.open(os.path.join(root, p1)).convert(\"RGB\")).unsqueeze(0).to(DEVICE)\n#             img2 = transform(Image.open(os.path.join(root, p2)).convert(\"RGB\")).unsqueeze(0).to(DEVICE)\n\n#             e1 = model(img1)\n#             e2 = model(img2)\n#             sims.append(F.cosine_similarity(e1, e2).item())\n#             labels.append(int(lab))\n\n#     fpr, tpr, thresholds = roc_curve(labels, sims)\n#     roc_auc = auc(fpr, tpr)\n#     eer = fpr[np.nanargmin(np.abs(fpr - (1 - tpr)))]\n#     best_thr = thresholds[np.nanargmin(np.abs(fpr - (1 - tpr)))]\n\n#     return roc_auc, eer, best_thr\n\ndef evaluate(model, pairs_file, root, transform):\n    model.eval()\n    sims, labels = [], []\n\n    with torch.no_grad():\n        for line in open(pairs_file):\n            parts = line.strip().split()\n            p1, p2, lab = parts[0], parts[1], parts[2]\n            \n            # Remove 'verification_data/' prefix if it exists\n            p1 = p1.replace('verification_data/', '')\n            p2 = p2.replace('verification_data/', '')\n            \n            img1_path = os.path.join(root, p1)\n            img2_path = os.path.join(root, p2)\n            \n            # Check if files exist\n            if not os.path.exists(img1_path) or not os.path.exists(img2_path):\n                print(f\"Warning: Skipping missing pair: {img1_path} or {img2_path}\")\n                continue\n            \n            img1 = transform(Image.open(img1_path).convert(\"RGB\")).unsqueeze(0).to(DEVICE)\n            img2 = transform(Image.open(img2_path).convert(\"RGB\")).unsqueeze(0).to(DEVICE)\n\n            e1 = model(img1)\n            e2 = model(img2)\n            sims.append(F.cosine_similarity(e1, e2).item())\n            labels.append(int(lab))\n\n    if len(sims) == 0:\n        print(\"ERROR: No valid pairs found!\")\n        return 0.0, 1.0, 0.5\n\n    fpr, tpr, thresholds = roc_curve(labels, sims)\n    roc_auc = auc(fpr, tpr)\n    eer = fpr[np.nanargmin(np.abs(fpr - (1 - tpr)))]\n    best_thr = thresholds[np.nanargmin(np.abs(fpr - (1 - tpr)))]\n\n    return roc_auc, eer, best_thr\n# ============================================================\n# TRAINING LOOP\n# ============================================================\n\ndef train():\n    transform = T.Compose([\n        T.Resize((IMG_SIZE, IMG_SIZE)),\n        T.RandomHorizontalFlip(),\n        T.ToTensor(),\n        T.Normalize([0.5]*3, [0.5]*3)\n    ])\n\n    train_ds = FaceFolderDataset(TRAIN_DATASET_PATH, transform)\n    val_ds = FaceFolderDataset(VAL_DATASET_PATH, transform)\n\n    sampler = BalancedIdentitySampler(\n        [l for _, l in train_ds.samples],\n        BATCH_IDENTITIES,\n        IMAGES_PER_ID\n    )\n\n    loader = DataLoader(\n        train_ds,\n        batch_size=BATCH_SIZE,\n        sampler=sampler,\n        num_workers=4,\n        pin_memory=True\n    )\n\n    model = FaceNet(EMBEDDING_DIM).to(DEVICE)\n    arcface = ArcFace(EMBEDDING_DIM, len(train_ds.labels), ARCFACE_MARGIN, ARCFACE_SCALE).to(DEVICE)\n\n    optimizer = torch.optim.AdamW(\n        list(model.parameters()) + list(arcface.parameters()),\n        lr=BASE_LR,\n        weight_decay=WEIGHT_DECAY\n    )\n\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n        optimizer,\n        T_max=EPOCHS_STAGE1 + EPOCHS_STAGE2\n    )\n\n    scaler = GradScaler()\n\n    best_auc = 0.0\n    best_state = None\n    best_threshold = 0.0\n\n    # ================= STAGE 1 =================\n    for epoch in range(EPOCHS_STAGE1):\n        model.train()\n        for imgs, labels in loader:\n            imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)\n\n            optimizer.zero_grad()\n            with autocast():\n                emb = model(imgs)\n                logits = arcface(emb, labels)\n                loss = F.cross_entropy(logits, labels)\n\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n\n        scheduler.step()\n\n        auc_val, eer, thr = evaluate(model, VAL_PAIRS_FILE, VAL_DATASET_PATH, transform)\n        if auc_val > best_auc:\n            best_auc = auc_val\n            best_state = model.state_dict()\n            best_threshold = thr\n\n    # ================= STAGE 2 =================\n    for epoch in range(EPOCHS_STAGE2):\n        model.train()\n        for imgs, labels in loader:\n            imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)\n\n            optimizer.zero_grad()\n            with autocast():\n                emb = model(imgs)\n                arc_loss = F.cross_entropy(arcface(emb, labels), labels)\n                tri_loss = triplet_loss(emb, labels)\n                loss = arc_loss + TRIPLET_WEIGHT * tri_loss\n\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n\n        scheduler.step()\n\n        auc_val, eer, thr = evaluate(model, VAL_PAIRS_FILE, VAL_DATASET_PATH, transform)\n        if auc_val > best_auc:\n            best_auc = auc_val\n            best_state = model.state_dict()\n            best_threshold = thr\n\n    # ========================================================\n    # SAVE BEST MODEL (ONCE)\n    # ========================================================\n\n    os.makedirs(OUTPUT_DIR, exist_ok=True)\n    torch.save(best_state, OUTPUT_MODEL_PATH)\n\n    metadata = {\n        \"embedding_dim\": EMBEDDING_DIM,\n        \"img_size\": IMG_SIZE,\n        \"threshold\": float(best_threshold),\n        \"best_auc\": float(best_auc),\n        \"model\": \"MobileNetV3 + ArcFace + Triplet\"\n    }\n\n    with open(OUTPUT_METADATA_PATH, \"w\") as f:\n        json.dump(metadata, f, indent=2)\n\n    print(\"Training complete. Best AUC:\", best_auc)\n\n\nif __name__ == \"__main__\":\n    train()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T06:42:38.897580Z","iopub.execute_input":"2026-01-15T06:42:38.898284Z","iopub.status.idle":"2026-01-15T06:43:44.984597Z","shell.execute_reply.started":"2026-01-15T06:42:38.898255Z","shell.execute_reply":"2026-01-15T06:43:44.983056Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"zip \n","metadata":{}},{"cell_type":"code","source":"!zip -r  arcfacetrain.zip /kaggle/working/","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}