{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":126777,"databundleVersionId":15314950}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -q timm umap-learn albumentations\ntry:\n    import faiss\nexcept ImportError:\n    !pip install -q faiss-cpu","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:34:45.563445Z","iopub.execute_input":"2026-03-12T23:34:45.563815Z","iopub.status.idle":"2026-03-12T23:34:54.932476Z","shell.execute_reply.started":"2026-03-12T23:34:45.563779Z","shell.execute_reply":"2026-03-12T23:34:54.931610Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport gc\nimport math\nimport json\nimport random\nimport warnings\nfrom pathlib import Path\nfrom collections import defaultdict\n\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\n\nimport cv2\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as mpatches\nimport seaborn as sns\n\nfrom sklearn.preprocessing import LabelEncoder\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import precision_score, recall_score, f1_score\nfrom sklearn.manifold import TSNE\n\nimport umap\nimport faiss\n\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers, regularizers\nfrom tensorflow.keras.applications import EfficientNetB3\nfrom tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, CSVLogger\n\nimport torch\nimport timm\n\nwarnings.filterwarnings('ignore')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:34:54.934423Z","iopub.execute_input":"2026-03-12T23:34:54.934766Z","iopub.status.idle":"2026-03-12T23:35:45.537136Z","shell.execute_reply.started":"2026-03-12T23:34:54.934738Z","shell.execute_reply":"2026-03-12T23:35:45.536030Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"SEED = 42\n\ndef set_seeds(seed=SEED):\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    random.seed(seed)\n    np.random.seed(seed)\n    tf.random.set_seed(seed)\n    torch.manual_seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)\n\nset_seeds()\n\ngpus = tf.config.list_physical_devices('GPU')\nfor gpu in gpus:\n    tf.config.experimental.set_memory_growth(gpu, True)\nprint(f\"GPUs: {[g.name for g in gpus]}\")\nprint(f\"TF: {tf.__version__}  PyTorch: {torch.__version__}  CUDA: {torch.cuda.is_available()}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:35:45.538801Z","iopub.execute_input":"2026-03-12T23:35:45.539576Z","iopub.status.idle":"2026-03-12T23:35:45.839000Z","shell.execute_reply.started":"2026-03-12T23:35:45.539550Z","shell.execute_reply":"2026-03-12T23:35:45.838183Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"CFG = {\n    'BASE_DIR'       : Path('/kaggle/input/competitions/jaguar-re-id'),\n    'OUTPUT_DIR'     : Path('/kaggle/working'),\n    'IMG_SIZE'       : 224,\n    'BATCH_SIZE'     : 32,\n    'EPOCHS'         : 50,\n    'LR'             : 1e-4,\n    'WEIGHT_DECAY'   : 1e-4,\n    'WARMUP_EPOCHS'  : 5,\n    'EMBEDDING_DIM'  : 512,\n    'DROPOUT'        : 0.3,\n    'ARC_MARGIN'     : 0.5,\n    'ARC_SCALE'      : 64,\n    'FOCAL_GAMMA'    : 2.0,\n    'FOCAL_ALPHA'    : 0.25,\n    'EARLY_STOP_PAT' : 10,\n    'VAL_SPLIT'      : 0.2,\n    'SEED'           : SEED,\n}\nCFG['TRAIN_DIR']  = CFG['BASE_DIR'] / 'train' / 'train'\nCFG['TEST_DIR']   = CFG['BASE_DIR'] / 'test'  / 'test'\nCFG['TRAIN_CSV']  = CFG['BASE_DIR'] / 'train.csv'\nCFG['TEST_CSV']   = CFG['BASE_DIR'] / 'test.csv'\nCFG['SUB_CSV']    = CFG['BASE_DIR'] / 'sample_submission.csv'\nCFG['OUTPUT_DIR'].mkdir(parents=True, exist_ok=True)\n\n# Verify paths exist\nfor key in ['TRAIN_DIR', 'TEST_DIR', 'TRAIN_CSV', 'TEST_CSV', 'SUB_CSV']:\n    exists = CFG[key].exists()\n    print(f\"  {key:<12} {'OK' if exists else 'MISSING'} — {CFG[key]}\")\n\nprint()\nprint(json.dumps({k: str(v) for k, v in CFG.items()}, indent=2))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-13T01:36:59.234841Z","iopub.execute_input":"2026-03-13T01:36:59.235673Z","iopub.status.idle":"2026-03-13T01:36:59.247761Z","shell.execute_reply.started":"2026-03-13T01:36:59.235645Z","shell.execute_reply":"2026-03-13T01:36:59.247122Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_csv = pd.read_csv(CFG['TRAIN_CSV'])\ntest_csv  = pd.read_csv(CFG['TEST_CSV'])\nsub_csv   = pd.read_csv(CFG['SUB_CSV'])\n\nprint(\"train.csv shape:\", train_csv.shape)\nprint(train_csv.head())\nprint(\"\\ntest.csv shape:\", test_csv.shape)\nprint(test_csv.head())\nprint(\"\\nsample_submission.csv shape:\", sub_csv.shape)\nprint(sub_csv.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:35:51.605481Z","iopub.execute_input":"2026-03-12T23:35:51.606189Z","iopub.status.idle":"2026-03-12T23:35:51.831962Z","shell.execute_reply.started":"2026-03-12T23:35:51.606131Z","shell.execute_reply":"2026-03-12T23:35:51.830966Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Attach full paths and encode labels\ntrain_csv['image_path'] = train_csv['filename'].apply(\n    lambda f: str(CFG['TRAIN_DIR'] / f)\n)\n\nle = LabelEncoder()\ntrain_csv['class_index'] = le.fit_transform(train_csv['ground_truth'])\n\nLABEL_ENCODER = le\nNUM_CLASSES   = int(train_csv['class_index'].nunique())\n\nprint(f\"Total train images : {len(train_csv)}\")\nprint(f\"Total identities   : {NUM_CLASSES}\")\nprint(f\"Classes sample     : {le.classes_[:10]}\")\ntrain_csv.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:35:53.767365Z","iopub.execute_input":"2026-03-12T23:35:53.767941Z","iopub.status.idle":"2026-03-12T23:35:53.808273Z","shell.execute_reply.started":"2026-03-12T23:35:53.767915Z","shell.execute_reply":"2026-03-12T23:35:53.807369Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class_counts = train_csv['ground_truth'].value_counts()\n\nprint(f\"Images            : {len(train_csv)}\")\nprint(f\"Identities        : {NUM_CLASSES}\")\nprint(f\"Mean imgs/identity: {class_counts.mean():.1f}\")\nprint(f\"Min               : {class_counts.min()}\")\nprint(f\"Max               : {class_counts.max()}\")\nprint(f\"Imbalance ratio   : {class_counts.max()/class_counts.min():.1f}x\")\n\nfig, axes = plt.subplots(1, 3, figsize=(20, 5))\nfig.suptitle('Jaguar Dataset — EDA', fontweight='bold')\n\naxes[0].bar(range(len(class_counts)), class_counts.values, color='steelblue', alpha=0.8)\naxes[0].axhline(class_counts.mean(), color='red', linestyle='--',\n                label=f'Mean={class_counts.mean():.1f}')\naxes[0].set_title('Class Frequency'); axes[0].set_xlabel('Identity'); axes[0].legend()\n\naxes[1].hist(class_counts.values, bins=30, color='coral', edgecolor='white')\naxes[1].set_title('Images per Identity Distribution')\n\nsorted_c = np.sort(class_counts.values)[::-1]\ncumsum   = np.cumsum(sorted_c) / sorted_c.sum() * 100\naxes[2].plot(range(1, len(cumsum)+1), cumsum, color='purple')\naxes[2].axhline(80, color='red', linestyle='--', label='80%')\naxes[2].set_title('Long-Tail CDF'); axes[2].legend()\n\nplt.tight_layout()\nplt.savefig(CFG['OUTPUT_DIR'] / 'eda.png', dpi=150)\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:35:56.600384Z","iopub.execute_input":"2026-03-12T23:35:56.601272Z","iopub.status.idle":"2026-03-12T23:35:57.584778Z","shell.execute_reply.started":"2026-03-12T23:35:56.601241Z","shell.execute_reply":"2026-03-12T23:35:57.584125Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"identities = train_csv['ground_truth'].unique()[:6]\nfig, axes  = plt.subplots(6, 4, figsize=(12, 18))\nfig.suptitle('Sample Images per Identity', fontweight='bold')\n\nfor row, identity in enumerate(identities):\n    paths = train_csv[train_csv['ground_truth'] == identity]['image_path'].tolist()\n    sample = random.sample(paths, min(4, len(paths)))\n    for col in range(4):\n        ax = axes[row][col]\n        ax.axis('off')\n        if col < len(sample):\n            img = cv2.cvtColor(cv2.imread(sample[col]), cv2.COLOR_BGR2RGB)\n            ax.imshow(img)\n        if col == 0:\n            ax.set_ylabel(identity, fontsize=9, rotation=0, labelpad=60)\n\nplt.tight_layout()\nplt.savefig(CFG['OUTPUT_DIR'] / 'sample_grid.png', dpi=150)\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:36:00.829128Z","iopub.execute_input":"2026-03-12T23:36:00.829714Z","iopub.status.idle":"2026-03-12T23:36:23.137272Z","shell.execute_reply.started":"2026-03-12T23:36:00.829683Z","shell.execute_reply":"2026-03-12T23:36:23.136403Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"IMG_SIZE = CFG['IMG_SIZE']\n\ndef load_and_preprocess(path, size=IMG_SIZE, crop_ratio=0.9):\n    img = cv2.imread(str(path))\n    if img is None:\n        return np.zeros((size, size, 3), dtype=np.float32)\n    img  = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    h, w = img.shape[:2]\n    ch, cw = int(h * crop_ratio), int(w * crop_ratio)\n    top, left = (h - ch) // 2, (w - cw) // 2\n    img = img[top:top+ch, left:left+cw]\n    img = cv2.resize(img, (size, size), interpolation=cv2.INTER_LANCZOS4)\n    # Keep in [0, 255] float32 — EfficientNetB3 has its own rescaling internally\n    return img.astype(np.float32)\n\n# QC\np = train_csv['image_path'].iloc[0]\nfig, axes = plt.subplots(1, 2, figsize=(8, 4))\naxes[0].imshow(cv2.cvtColor(cv2.imread(p), cv2.COLOR_BGR2RGB))\naxes[0].set_title('Original'); axes[0].axis('off')\naxes[1].imshow(load_and_preprocess(p).astype(np.uint8))\naxes[1].set_title('Preprocessed'); axes[1].axis('off')\nplt.tight_layout(); plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:36:23.138731Z","iopub.execute_input":"2026-03-12T23:36:23.139144Z","iopub.status.idle":"2026-03-12T23:36:23.836430Z","shell.execute_reply.started":"2026-03-12T23:36:23.139098Z","shell.execute_reply":"2026-03-12T23:36:23.835526Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def build_augmentation_layer():\n    return keras.Sequential([\n        layers.RandomFlip('horizontal'),\n        layers.RandomRotation(factor=0.08),\n        layers.RandomZoom(height_factor=(-0.1, 0.1)),\n        layers.RandomTranslation(height_factor=0.05, width_factor=0.05),\n    ], name='augmentation')\n# Brightness/Contrast/Noise removed — they are unstable on [0,255] float\n# and cause black frames. Use albumentations for those if needed later.\n\naug    = build_augmentation_layer()\norigin = load_and_preprocess(train_csv['image_path'].iloc[0])  # [0,255]\n\nfig, axes = plt.subplots(2, 5, figsize=(18, 7))\nfig.suptitle('Augmentation Samples', fontweight='bold')\n\naxes[0][0].imshow(origin.astype(np.uint8))\naxes[0][0].set_title('Original')\naxes[0][0].axis('off')\n\nfor i in range(1, 10):\n    r, c = divmod(i, 5)\n    t    = tf.expand_dims(tf.constant(origin), 0)\n    out  = aug(t, training=True)[0].numpy()\n    out  = np.clip(out, 0, 255).astype(np.uint8)\n    axes[r][c].imshow(out)\n    axes[r][c].set_title(f'Aug {i}')\n    axes[r][c].axis('off')\n\nplt.tight_layout()\nplt.savefig(CFG['OUTPUT_DIR'] / 'augmentation_samples.png', dpi=150)\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:36:23.837552Z","iopub.execute_input":"2026-03-12T23:36:23.837872Z","iopub.status.idle":"2026-03-12T23:36:29.389866Z","shell.execute_reply.started":"2026-03-12T23:36:23.837849Z","shell.execute_reply":"2026-03-12T23:36:29.388998Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_split, val_split = train_test_split(\n    train_csv,\n    test_size=CFG['VAL_SPLIT'],\n    stratify=train_csv['class_index'],\n    random_state=SEED\n)\ntrain_split = train_split.reset_index(drop=True)\nval_split   = val_split.reset_index(drop=True)\nprint(f\"Train: {len(train_split)}  Val: {len(val_split)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:36:29.392032Z","iopub.execute_input":"2026-03-12T23:36:29.392386Z","iopub.status.idle":"2026-03-12T23:36:29.403835Z","shell.execute_reply.started":"2026-03-12T23:36:29.392366Z","shell.execute_reply":"2026-03-12T23:36:29.403235Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"AUTOTUNE = tf.data.AUTOTUNE\n\ndef parse_image_tf(path, label):\n    raw = tf.io.read_file(path)\n    img = tf.image.decode_image(raw, channels=3, expand_animations=False)\n    img = tf.image.resize(img, [IMG_SIZE, IMG_SIZE], method='lanczos3')\n    img = tf.cast(img, tf.float32)\n    # Do NOT divide by 255 — EfficientNetB3 preprocesses internally\n    return img, label\n\ndef build_dataset(df, augment=True, shuffle=True):\n    paths  = df['image_path'].tolist()\n    labels = df['class_index'].tolist()\n    ds     = tf.data.Dataset.from_tensor_slices((paths, labels))\n    if shuffle:\n        ds = ds.shuffle(len(paths), seed=SEED, reshuffle_each_iteration=True)\n    ds = ds.map(parse_image_tf, num_parallel_calls=AUTOTUNE).cache()\n    if augment:\n        aug_l = build_augmentation_layer()\n        ds    = ds.map(\n            lambda x, y: (\n                tf.clip_by_value(\n                    aug_l(tf.expand_dims(x, 0), training=True)[0], 0, 255\n                ), y\n            ),\n            num_parallel_calls=AUTOTUNE\n        )\n    return ds.batch(CFG['BATCH_SIZE'], drop_remainder=False).prefetch(AUTOTUNE)\n\ntrain_ds = build_dataset(train_split, augment=True,  shuffle=True)\nval_ds   = build_dataset(val_split,   augment=False, shuffle=False)\n\nfor imgs, lbls in train_ds.take(1):\n    print(f\"Batch: {imgs.shape}  Labels: {lbls.shape}  Range: [{imgs.numpy().min():.1f}, {imgs.numpy().max():.1f}]\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:36:29.404792Z","iopub.execute_input":"2026-03-12T23:36:29.405111Z","iopub.status.idle":"2026-03-12T23:36:35.829985Z","shell.execute_reply.started":"2026-03-12T23:36:29.405080Z","shell.execute_reply":"2026-03-12T23:36:35.829350Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class FocalLoss(keras.losses.Loss):\n    def __init__(self, gamma=CFG['FOCAL_GAMMA'], alpha=CFG['FOCAL_ALPHA'], **kwargs):\n        super().__init__(**kwargs)\n        self.gamma = gamma\n        self.alpha = alpha\n\n    def call(self, y_true, y_pred):\n        p     = tf.clip_by_value(tf.nn.softmax(y_pred), 1e-7, 1.0)\n        oh    = tf.one_hot(tf.cast(y_true, tf.int32), tf.shape(y_pred)[-1])\n        p_t   = tf.reduce_sum(oh * p, axis=-1)\n        return tf.reduce_mean(self.alpha * tf.pow(1 - p_t, self.gamma) * (-tf.math.log(p_t)))\n\n    def get_config(self):\n        return {**super().get_config(), 'gamma': self.gamma, 'alpha': self.alpha}\n\n\ndef compute_class_weights(df):\n    counts = df['class_index'].value_counts().sort_index()\n    total  = len(df)\n    n      = len(counts)\n    return {idx: total / (n * cnt) for idx, cnt in counts.items()}\n\nclass_weights = compute_class_weights(train_split)\nprint(f\"Weight range: [{min(class_weights.values()):.4f}, {max(class_weights.values()):.4f}]\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:36:35.831291Z","iopub.execute_input":"2026-03-12T23:36:35.831502Z","iopub.status.idle":"2026-03-12T23:36:35.844765Z","shell.execute_reply.started":"2026-03-12T23:36:35.831484Z","shell.execute_reply":"2026-03-12T23:36:35.843934Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ArcFaceLayer(layers.Layer):\n    def __init__(self, num_classes, embedding_dim,\n                 margin=0.3,   # reduced from 0.5 — less aggressive early on\n                 scale=32,     # reduced from 64 — softer logit scaling\n                 **kwargs):\n        super().__init__(**kwargs)\n        self.num_classes   = num_classes\n        self.embedding_dim = embedding_dim\n        self.margin        = margin\n        self.scale         = scale\n        self.cos_m = float(math.cos(margin))\n        self.sin_m = float(math.sin(margin))\n        self.th    = float(math.cos(math.pi - margin))\n        self.mm    = float(math.sin(math.pi - margin) * margin)\n\n    def build(self, input_shape):\n        self.W = self.add_weight(\n            shape=(self.embedding_dim, self.num_classes),\n            initializer='glorot_uniform', trainable=True, name='W'\n        )\n\n    def call(self, embeddings, labels=None, training=None):\n        emb = tf.math.l2_normalize(embeddings, axis=1)\n        W   = tf.math.l2_normalize(self.W, axis=0)\n        cos = tf.clip_by_value(tf.matmul(emb, W), -1+1e-7, 1-1e-7)\n        if labels is None or not training:\n            return self.scale * cos\n        sin     = tf.sqrt(tf.maximum(1.0 - cos**2, 1e-9))\n        phi     = cos * self.cos_m - sin * self.sin_m\n        phi     = tf.where(cos > self.th, phi, cos - self.mm)\n        oh      = tf.one_hot(tf.cast(labels, tf.int32), self.num_classes)\n        return self.scale * tf.where(tf.cast(oh, tf.bool), phi, cos)\n\n    def get_config(self):\n        return {**super().get_config(), 'num_classes': self.num_classes,\n                'embedding_dim': self.embedding_dim,\n                'margin': self.margin, 'scale': self.scale}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:36:35.845737Z","iopub.execute_input":"2026-03-12T23:36:35.846001Z","iopub.status.idle":"2026-03-12T23:36:35.862510Z","shell.execute_reply.started":"2026-03-12T23:36:35.845975Z","shell.execute_reply":"2026-03-12T23:36:35.861642Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def build_efficientnet_model(num_classes, emb_dim=CFG['EMBEDDING_DIM'],\n                              dropout=CFG['DROPOUT'], img_size=IMG_SIZE):\n    img_in  = keras.Input(shape=(img_size, img_size, 3), name='img')\n    lbl_in  = keras.Input(shape=(), dtype=tf.int32,       name='lbl')\n\n    x        = build_augmentation_layer()(img_in)\n    backbone = EfficientNetB3(include_top=False, weights='imagenet',\n                               input_tensor=x, pooling=None)\n\n    # Unfreeze entire backbone — fine-tune everything with weight decay\n    for layer in backbone.layers:\n        layer.trainable = True\n\n    x = backbone.output\n    x = layers.GlobalAveragePooling2D()(x)\n    x = layers.BatchNormalization()(x)\n    x = layers.Dropout(dropout)(x)\n    x = layers.Dense(emb_dim, use_bias=False,\n                      kernel_regularizer=regularizers.l2(1e-4),\n                      name='embedding')(x)\n    x = layers.BatchNormalization(name='emb_bn')(x)\n    embeddings = layers.Lambda(\n        lambda t: tf.math.l2_normalize(t, axis=1), name='l2_norm'\n    )(x)\n    logits = ArcFaceLayer(num_classes, emb_dim, name='arcface')(embeddings, lbl_in)\n\n    model = keras.Model(inputs=[img_in, lbl_in], outputs=logits, name='EfficientNet_ReID')\n    model.embedding_model = keras.Model(\n        inputs=img_in, outputs=embeddings, name='EfficientNet_Emb'\n    )\n    return model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:36:35.863672Z","iopub.execute_input":"2026-03-12T23:36:35.864292Z","iopub.status.idle":"2026-03-12T23:36:35.877252Z","shell.execute_reply.started":"2026-03-12T23:36:35.864270Z","shell.execute_reply":"2026-03-12T23:36:35.876462Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ArcFaceTrainer(keras.Model):\n    def __init__(self, model, loss_fn):\n        super().__init__()\n        self.arc_model  = model\n        self.loss_fn    = loss_fn\n        self.acc_metric = keras.metrics.SparseCategoricalAccuracy(name='acc')\n        self.loss_metric = keras.metrics.Mean(name='loss')\n\n    def train_step(self, data):\n        if len(data) == 3:\n            images, labels, _ = data\n        else:\n            images, labels = data\n\n        with tf.GradientTape() as tape:\n            logits = self.arc_model([images, labels], training=True)\n            loss   = self.loss_fn(labels, logits) + sum(self.arc_model.losses)\n\n        grads = tape.gradient(loss, self.arc_model.trainable_variables)\n        self.optimizer.apply_gradients(zip(grads, self.arc_model.trainable_variables))\n        self.loss_metric.update_state(loss)\n        self.acc_metric.update_state(labels, logits)\n        return {'loss': self.loss_metric.result(),\n                'acc':  self.acc_metric.result()}\n\n    def test_step(self, data):\n        if len(data) == 3:\n            images, labels, _ = data\n        else:\n            images, labels = data\n\n        logits = self.arc_model([images, tf.zeros_like(labels)], training=False)\n        loss   = self.loss_fn(labels, logits)\n        self.loss_metric.update_state(loss)\n        self.acc_metric.update_state(labels, logits)\n        return {'loss': self.loss_metric.result(),\n                'acc':  self.acc_metric.result()}\n\n    @property\n    def metrics(self):\n        # returning metrics here ensures Keras resets them each epoch\n        return [self.loss_metric, self.acc_metric]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:36:35.878249Z","iopub.execute_input":"2026-03-12T23:36:35.878790Z","iopub.status.idle":"2026-03-12T23:36:35.895620Z","shell.execute_reply.started":"2026-03-12T23:36:35.878769Z","shell.execute_reply":"2026-03-12T23:36:35.894961Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"steps_per_epoch = math.ceil(len(train_split) / CFG['BATCH_SIZE'])\ntotal_steps     = steps_per_epoch * CFG['EPOCHS']\nwarmup_steps    = steps_per_epoch * 2  # only 2 epochs warmup\n\nclass WarmupCosine(keras.optimizers.schedules.LearningRateSchedule):\n    def __init__(self, base_lr, total_steps, warmup_steps, min_lr=1e-6):\n        self.base_lr      = base_lr\n        self.total_steps  = total_steps\n        self.warmup_steps = warmup_steps\n        self.min_lr       = min_lr\n\n    def __call__(self, step):\n        step    = tf.cast(step, tf.float32)\n        warmup  = self.base_lr * step / tf.maximum(\n            tf.cast(self.warmup_steps, tf.float32), 1.0\n        )\n        decay   = self.min_lr + 0.5 * (self.base_lr - self.min_lr) * (\n            1 + tf.cos(math.pi * (step - self.warmup_steps) /\n                       tf.maximum(\n                           tf.cast(self.total_steps - self.warmup_steps, tf.float32), 1.0\n                       ))\n        )\n        return tf.where(step < self.warmup_steps, warmup, decay)\n\n    def get_config(self):\n        return {'base_lr': self.base_lr, 'total_steps': self.total_steps,\n                'warmup_steps': self.warmup_steps}\n\nlr_schedule = WarmupCosine(\n    base_lr=3e-4,   # increased from 1e-4\n    total_steps=total_steps,\n    warmup_steps=warmup_steps\n)\nprint(f\"Steps/epoch:{steps_per_epoch}  Total:{total_steps}  Warmup:{warmup_steps}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:36:35.897675Z","iopub.execute_input":"2026-03-12T23:36:35.897893Z","iopub.status.idle":"2026-03-12T23:36:35.912344Z","shell.execute_reply.started":"2026-03-12T23:36:35.897875Z","shell.execute_reply":"2026-03-12T23:36:35.911741Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"efficientnet_model = build_efficientnet_model(NUM_CLASSES)\neff_trainer        = ArcFaceTrainer(efficientnet_model, FocalLoss())\n\neff_trainer.compile(optimizer=keras.optimizers.AdamW(\n    learning_rate=lr_schedule, weight_decay=CFG['WEIGHT_DECAY']\n))\n\nclass InnerModelCheckpoint(keras.callbacks.Callback):\n    def __init__(self, filepath, monitor='val_acc', mode='max'):\n        super().__init__()\n        self.filepath = filepath\n        self.monitor  = monitor\n        self.mode     = mode\n        self.best     = -np.inf if mode == 'max' else np.inf\n\n    def on_epoch_end(self, epoch, logs=None):\n        current = logs.get(self.monitor)\n        if current is None:\n            return\n        improved = current > self.best if self.mode == 'max' else current < self.best\n        if improved:\n            self.best = current\n            self.model.arc_model.save(self.filepath)\n            size_mb = Path(self.filepath).stat().st_size / 1024**2\n            print(f\"\\nCheckpoint saved ({size_mb:.1f} MB) — \"\n                  f\"{self.monitor}: {current:.4f}\")\n            if size_mb < 1.0:\n                print(\"WARNING: file too small — save may have failed\")\n\neff_callbacks = [\n    EarlyStopping(monitor='val_loss', patience=CFG['EARLY_STOP_PAT'],\n                  restore_best_weights=False, verbose=1),\n    InnerModelCheckpoint(\n        str(CFG['OUTPUT_DIR'] / 'eff_best.keras'),\n        monitor='val_acc', mode='max'\n    ),\n    CSVLogger(str(CFG['OUTPUT_DIR'] / 'eff_history.csv'))\n]\n\neff_history = eff_trainer.fit(\n    train_ds, validation_data=val_ds,\n    epochs=CFG['EPOCHS'], callbacks=eff_callbacks,\n    verbose=1\n)\n\n# Force save after training as backup\nefficientnet_model.save(str(CFG['OUTPUT_DIR'] / 'efficientnet_final.keras'))\nefficientnet_model.embedding_model.save(\n    str(CFG['OUTPUT_DIR'] / 'efficientnet_embedding.keras')\n)\n\n# Verify\nfor fname in ['eff_best.keras', 'efficientnet_final.keras', 'efficientnet_embedding.keras']:\n    path = CFG['OUTPUT_DIR'] / fname\n    if path.exists():\n        print(f\"{fname}: {path.stat().st_size/1024**2:.2f} MB\")\n    else:\n        print(f\"ERROR: {fname} not found\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:36:35.913319Z","iopub.execute_input":"2026-03-12T23:36:35.913579Z","iopub.status.idle":"2026-03-12T23:50:27.326494Z","shell.execute_reply.started":"2026-03-12T23:36:35.913552Z","shell.execute_reply":"2026-03-12T23:50:27.325813Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ── Save EfficientNet after training ────────────────────────────────────────\neff_save_path = CFG['OUTPUT_DIR'] / 'efficientnet_final.weights.h5'\nefficientnet_model.save_weights(str(eff_save_path))\nprint(f\"EfficientNet saved → {eff_save_path}\")\n\n# Also save full embedding model separately\neff_emb_save_path = CFG['OUTPUT_DIR'] / 'efficientnet_embedding.weights.h5'\nefficientnet_model.embedding_model.save_weights(str(eff_emb_save_path))\nprint(f\"EfficientNet embedding model saved → {eff_emb_save_path}\")\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:52:00.662350Z","iopub.execute_input":"2026-03-12T23:52:00.662840Z","iopub.status.idle":"2026-03-12T23:52:02.301353Z","shell.execute_reply.started":"2026-03-12T23:52:00.662814Z","shell.execute_reply":"2026-03-12T23:52:02.300572Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import gc\ntf.keras.backend.clear_session()\ngc.collect()\nif torch.cuda.is_available():\n    torch.cuda.empty_cache()\n    print(f\"GPU free: {torch.cuda.mem_get_info()[0]/1024**3:.2f} GB\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:52:22.183574Z","iopub.execute_input":"2026-03-12T23:52:22.184143Z","iopub.status.idle":"2026-03-12T23:52:24.585519Z","shell.execute_reply.started":"2026-03-12T23:52:22.184109Z","shell.execute_reply":"2026-03-12T23:52:24.584808Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class SwinExtractor:\n    def __init__(self,\n                 model_name='swin_small_patch4_window7_224',\n                 emb_dim=CFG['EMBEDDING_DIM'], num_classes=NUM_CLASSES):\n        self.device     = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n        self.emb_dim    = emb_dim\n        self.arc_margin = 0.3\n        self.arc_scale  = 32\n\n        self.backbone = timm.create_model(\n            model_name,\n            pretrained=True,\n            num_classes=0\n        ).to(self.device)\n\n        if hasattr(self.backbone, 'set_grad_checkpointing'):\n            self.backbone.set_grad_checkpointing(True)\n            print(\"Gradient checkpointing enabled\")\n\n        in_feat = self.backbone.num_features\n\n        self.projector = torch.nn.Sequential(\n            torch.nn.Linear(in_feat, emb_dim, bias=False),\n            torch.nn.BatchNorm1d(emb_dim)\n        ).to(self.device)\n\n        self.head = torch.nn.Linear(emb_dim, num_classes, bias=False).to(self.device)\n        torch.nn.init.normal_(self.head.weight, std=0.01)\n\n        self.params = (\n            list(self.backbone.parameters()) +\n            list(self.projector.parameters()) +\n            list(self.head.parameters())\n        )\n\n        total = sum(p.numel() for p in self.backbone.parameters())\n        print(f\"Swin-Small | device:{self.device} | params:{total/1e6:.1f}M\")\n\n    def _norm_batch(self, images):\n        MEAN = torch.tensor([0.485, 0.456, 0.406], device=self.device)\n        STD  = torch.tensor([0.229, 0.224, 0.225], device=self.device)\n        t    = torch.from_numpy(images).permute(0, 3, 1, 2).float().to(self.device)\n        return (t - MEAN[None,:,None,None]) / STD[None,:,None,None]\n\n    def forward(self, images, labels=None, training=False):\n        self.backbone.train(training)\n        self.projector.train(training)\n        self.head.train(training)\n\n        with torch.set_grad_enabled(training):\n            feat = self.backbone(self._norm_batch(images))\n            emb  = torch.nn.functional.normalize(self.projector(feat), dim=1)\n\n            W   = torch.nn.functional.normalize(self.head.weight, dim=1)\n            cos = torch.matmul(emb, W.T).clamp(-1+1e-7, 1-1e-7)\n\n            if training and labels is not None:\n                sin = torch.sqrt(1 - cos**2 + 1e-9)\n                phi = (cos * math.cos(self.arc_margin)\n                       - sin * math.sin(self.arc_margin))\n                lbl = torch.from_numpy(labels).long().to(self.device)\n                oh  = torch.zeros_like(cos).scatter_(1, lbl.unsqueeze(1), 1.0)\n                out = torch.where(oh.bool(), phi, cos)\n                return emb.detach().cpu().numpy(), self.arc_scale * out\n\n            return emb.detach().cpu().numpy(), cos\n\n    def save(self, path):\n        torch.save({\n            'backbone' : self.backbone.state_dict(),\n            'projector': self.projector.state_dict(),\n            'head'     : self.head.state_dict()\n        }, path)\n        print(f\"Saved → {path}\")\n\n    def load(self, path):\n        ck = torch.load(path, map_location=self.device)\n        self.backbone.load_state_dict(ck['backbone'])\n        self.projector.load_state_dict(ck['projector'])\n        self.head.load_state_dict(ck['head'])\n        print(f\"Loaded ← {path}\")\n\nswin = SwinExtractor()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:52:32.290735Z","iopub.execute_input":"2026-03-12T23:52:32.291238Z","iopub.status.idle":"2026-03-12T23:52:35.084117Z","shell.execute_reply.started":"2026-03-12T23:52:32.291200Z","shell.execute_reply":"2026-03-12T23:52:35.083107Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def preload_images(df, desc='Loading'):\n    \"\"\"Load all images into RAM once — eliminates per-batch disk I/O.\"\"\"\n    images = {}\n    for path in tqdm(df['image_path'].tolist(), desc=desc):\n        images[path] = load_and_preprocess(path)\n    return images\n\ndef train_swin(swin, df_train, df_val,\n               epochs=CFG['EPOCHS'],\n               bs=16,\n               lr=CFG['LR']):\n    opt      = torch.optim.AdamW(swin.params, lr=lr, weight_decay=CFG['WEIGHT_DECAY'])\n    crit     = torch.nn.CrossEntropyLoss()\n    sched    = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs, eta_min=1e-6)\n    aug      = build_augmentation_layer()\n    history  = defaultdict(list)\n    best_acc = 0.0\n    patience = 0\n\n    # ── Preload all images into RAM ──────────────────────────────────────\n    print(\"Preloading train images into RAM...\")\n    train_cache = preload_images(df_train, desc='Train images')\n    print(\"Preloading val images into RAM...\")\n    val_cache   = preload_images(df_val,   desc='Val images')\n    print(f\"Cached {len(train_cache)} train + {len(val_cache)} val images\")\n\n    # Pre-stack val images once — never changes\n    val_imgs   = np.stack([val_cache[p] for p in df_val['image_path']]) / 255.0\n    val_labels_arr = df_val['class_index'].values.astype(np.int64)\n\n    for epoch in range(epochs):\n        swin.backbone.train(); swin.projector.train(); swin.head.train()\n        idx = np.random.permutation(len(df_train))\n        tr_loss = tr_corr = tr_tot = 0\n\n        for b in range(math.ceil(len(df_train)/bs)):\n            batch  = df_train.iloc[idx[b*bs:(b+1)*bs]]\n            imgs   = np.stack([train_cache[p] for p in batch['image_path']])\n            labels = batch['class_index'].values.astype(np.int64)\n\n            # Augment in [0,255], normalize to [0,1]\n            imgs = np.clip(aug(tf.constant(imgs), training=True).numpy(), 0, 255)\n            imgs = imgs / 255.0\n\n            opt.zero_grad()\n            _, logits = swin.forward(imgs, labels, training=True)\n            loss = crit(logits, torch.from_numpy(labels).long().to(swin.device))\n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(swin.params, 1.0)\n            opt.step()\n            torch.cuda.empty_cache()\n\n            tr_loss += loss.item()\n            tr_corr += (logits.detach().cpu().numpy().argmax(1) == labels).sum()\n            tr_tot  += len(labels)\n\n        sched.step()\n        swin.backbone.eval(); swin.projector.eval(); swin.head.eval()\n        vl_loss = vl_corr = vl_tot = 0\n\n        # Val in batches from pre-stacked array\n        for b in range(math.ceil(len(val_imgs)/bs)):\n            imgs_b   = val_imgs[b*bs:(b+1)*bs]\n            labels_b = val_labels_arr[b*bs:(b+1)*bs]\n            with torch.no_grad():\n                emb, lv = swin.forward(imgs_b, training=False)\n                lt       = torch.from_numpy(labels_b).long().to(swin.device)\n                vl_loss += crit(lv, lt).item()\n                vl_corr += (lv.cpu().numpy().argmax(1) == labels_b).sum()\n                vl_tot  += len(labels_b)\n\n        n_tr = math.ceil(len(df_train)/bs)\n        n_vl = math.ceil(len(val_imgs)/bs)\n        ep   = {\n            'loss'    : tr_loss / n_tr,\n            'acc'     : tr_corr / tr_tot,\n            'val_loss': vl_loss / n_vl,\n            'val_acc' : vl_corr / vl_tot\n        }\n        for k, v in ep.items():\n            history[k].append(float(v))\n\n        print(f\"Epoch {epoch+1:03d}  \"\n              f\"loss:{ep['loss']:.4f} acc:{ep['acc']:.4f}  \"\n              f\"val_loss:{ep['val_loss']:.4f} val_acc:{ep['val_acc']:.4f}\")\n\n        if ep['val_acc'] > best_acc:\n            best_acc = ep['val_acc']\n            swin.save(str(CFG['OUTPUT_DIR'] / 'swin_best.pt'))\n            patience = 0\n        else:\n            patience += 1\n            if patience >= CFG['EARLY_STOP_PAT']:\n                print(f\"Early stop at epoch {epoch+1}\")\n                break\n\n    return dict(history)\n\nswin_history = train_swin(swin, train_split, val_split)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T23:52:41.577964Z","iopub.execute_input":"2026-03-12T23:52:41.578911Z","iopub.status.idle":"2026-03-13T00:38:17.968152Z","shell.execute_reply.started":"2026-03-12T23:52:41.578884Z","shell.execute_reply":"2026-03-13T00:38:17.967317Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Save and verify Swin\nswin_final_path = CFG['OUTPUT_DIR'] / 'swin_final.pt'\nswin_best_path  = CFG['OUTPUT_DIR'] / 'swin_best.pt'\n\nswin.save(str(swin_final_path))\n\n# If best was never saved during training, save final as best\nif not swin_best_path.exists():\n    swin.save(str(swin_best_path))\n\n# Verify both\nfor path in [swin_final_path, swin_best_path]:\n    size_mb = path.stat().st_size / 1024**2\n    print(f\"{path.name}: {size_mb:.2f} MB\")\n    if size_mb < 10:\n        print(f\"ERROR: {path.name} too small — save failed\")\n\n# List all outputs\nprint(\"\\nAll saved files:\")\nfor f in sorted(CFG['OUTPUT_DIR'].glob('*')):\n    print(f\"  {f.name:<45} {f.stat().st_size/1024**2:.2f} MB\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-13T00:38:17.969983Z","iopub.execute_input":"2026-03-13T00:38:17.970342Z","iopub.status.idle":"2026-03-13T00:38:18.206621Z","shell.execute_reply.started":"2026-03-13T00:38:17.970320Z","shell.execute_reply":"2026-03-13T00:38:18.205883Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for f in sorted(CFG['OUTPUT_DIR'].glob('*')):\n    print(f\"  {f.name:<45} {f.stat().st_size/1024**2:.2f} MB\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-13T01:09:00.385891Z","iopub.execute_input":"2026-03-13T01:09:00.386620Z","iopub.status.idle":"2026-03-13T01:09:00.392396Z","shell.execute_reply.started":"2026-03-13T01:09:00.386594Z","shell.execute_reply":"2026-03-13T01:09:00.391601Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ── Step 1: Define L2NormLayer ────────────────────────────────────────────────\nclass L2NormLayer(layers.Layer):\n    def call(self, x):\n        return tf.math.l2_normalize(x, axis=1)\n    def get_config(self):\n        return super().get_config()\n\n# ── Step 2: Rebuild model and load weights from .h5 ──────────────────────────\nefficientnet_model = build_efficientnet_model(NUM_CLASSES)\n\n# Dummy pass to build all layers before loading weights\ndummy_imgs   = np.zeros((1, IMG_SIZE, IMG_SIZE, 3), dtype=np.float32)\ndummy_labels = np.zeros((1,), dtype=np.int32)\n_ = efficientnet_model([dummy_imgs, dummy_labels], training=False)\n\n# Load weights — .weights.h5 has no architecture baked in, just raw weights\nefficientnet_model.load_weights(\n    str(CFG['OUTPUT_DIR'] / 'efficientnet_final.weights.h5')\n)\nprint(f\"EfficientNet weights loaded — \"\n      f\"{(CFG['OUTPUT_DIR'] / 'efficientnet_final.weights.h5').stat().st_size/1024**2:.2f} MB\")\n\neff_emb_model = efficientnet_model.embedding_model\n\n# ── Step 3: Load Swin ─────────────────────────────────────────────────────────\nswin = SwinExtractor()\nswin.load(str(CFG['OUTPUT_DIR'] / 'swin_best.pt'))\nprint(f\"Swin loaded — \"\n      f\"{(CFG['OUTPUT_DIR'] / 'swin_best.pt').stat().st_size/1024**2:.2f} MB\")\n\n# ── Step 4: Verify ────────────────────────────────────────────────────────────\nprint(\"\\nVerifying embedding quality...\")\nsample_imgs = np.stack([\n    load_and_preprocess(p)\n    for p in train_split['image_path'].tolist()[:4]\n])\n\neff_sample = eff_emb_model.predict(sample_imgs, verbose=0)\nprint(f\"EfficientNet std : {eff_sample.std(axis=0).mean():.6f}\")\nprint(f\"EfficientNet same: {np.allclose(eff_sample[0], eff_sample[1])}\")\n\nswin_sample, _ = swin.forward(sample_imgs / 255.0, training=False)\nprint(f\"Swin std         : {swin_sample.std(axis=0).mean():.6f}\")\nprint(f\"Swin same        : {np.allclose(swin_sample[0], swin_sample[1])}\")\n\nif np.allclose(eff_sample[0], eff_sample[1]):\n    raise ValueError(\"EfficientNet embeddings identical — weights not loaded correctly\")\nif np.allclose(swin_sample[0], swin_sample[1]):\n    raise ValueError(\"Swin embeddings identical — weights not loaded correctly\")\n\nprint(\"\\nBoth models verified OK\")\n\n# ── Step 5: Extraction functions ──────────────────────────────────────────────\ndef extract_eff(df, bs=64):\n    all_embs = []\n    for b in tqdm(range(math.ceil(len(df)/bs)), desc='EfficientNet'):\n        imgs = np.stack([load_and_preprocess(p)\n                         for p in df.iloc[b*bs:(b+1)*bs]['image_path']])\n        all_embs.append(eff_emb_model.predict(imgs, verbose=0))\n    return np.vstack(all_embs)\n\ndef extract_swin(df, bs=32):\n    all_embs = []\n    for b in tqdm(range(math.ceil(len(df)/bs)), desc='Swin'):\n        imgs = np.stack([load_and_preprocess(p)\n                         for p in df.iloc[b*bs:(b+1)*bs]['image_path']])\n        embs, _ = swin.forward(imgs / 255.0, training=False)\n        all_embs.append(embs)\n    return np.vstack(all_embs)\n\n# ── Step 6: Extract embeddings ────────────────────────────────────────────────\nprint(\"\\nExtracting train embeddings (gallery)...\")\neff_train_embs  = extract_eff(train_split)\nswin_train_embs = extract_swin(train_split)\n\nprint(\"\\nExtracting val embeddings (queries)...\")\neff_val_embs  = extract_eff(val_split)\nswin_val_embs = extract_swin(val_split)\n\ntrain_labels = train_split['class_index'].values\nval_labels   = val_split['class_index'].values\n\nprint(f\"\\nEfficientNet embs : {eff_val_embs.shape}\")\nprint(f\"Swin embs         : {swin_val_embs.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-13T01:09:39.608331Z","iopub.execute_input":"2026-03-13T01:09:39.608958Z","iopub.status.idle":"2026-03-13T01:28:51.730428Z","shell.execute_reply.started":"2026-03-13T01:09:39.608930Z","shell.execute_reply":"2026-03-13T01:28:51.729743Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fig, axes = plt.subplots(2, 2, figsize=(14, 8))\nfig.suptitle('Training History', fontweight='bold')\n\naxes[0][0].plot(eff_history.history['loss'],     label='Train')\naxes[0][0].plot(eff_history.history['val_loss'], label='Val')\naxes[0][0].set_title('EfficientNet Loss'); axes[0][0].legend()\n\naxes[0][1].plot(eff_history.history['acc'],      label='Train')\naxes[0][1].plot(eff_history.history['val_acc'],  label='Val')\naxes[0][1].set_title('EfficientNet Accuracy'); axes[0][1].legend()\n\naxes[1][0].plot(swin_history['loss'],            label='Train')\naxes[1][0].plot(swin_history['val_loss'],        label='Val')\naxes[1][0].set_title('Swin Loss'); axes[1][0].legend()\n\naxes[1][1].plot(swin_history['acc'],             label='Train')\naxes[1][1].plot(swin_history['val_acc'],         label='Val')\naxes[1][1].set_title('Swin Accuracy'); axes[1][1].legend()\n\nplt.tight_layout()\nplt.savefig(CFG['OUTPUT_DIR'] / 'training_curves.png', dpi=150)\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-13T01:29:21.869944Z","iopub.execute_input":"2026-03-13T01:29:21.870406Z","iopub.status.idle":"2026-03-13T01:29:22.926894Z","shell.execute_reply.started":"2026-03-13T01:29:21.870385Z","shell.execute_reply":"2026-03-13T01:29:22.926252Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def build_faiss_index(embs):\n    e = embs.astype(np.float32)\n    faiss.normalize_L2(e)\n    idx = faiss.IndexFlatL2(e.shape[1])\n    idx.add(e)\n    return idx\n\ndef retrieval_metrics(val_embs, train_embs, val_labels, train_labels, name):\n    idx   = build_faiss_index(train_embs)\n    q     = val_embs.astype(np.float32); faiss.normalize_L2(q)\n    _, nn = idx.search(q, 5)\n    top1  = train_labels[nn[:, 0]]\n    top5  = train_labels[nn[:, :5]]\n    top1_acc = (top1 == val_labels).mean()\n    top5_acc = sum(val_labels[i] in top5[i] for i in range(len(val_labels))) / len(val_labels)\n    return {\n        'Model'     : name,\n        'Top-1 Acc' : top1_acc,\n        'Top-5 Acc' : top5_acc,\n        'Precision' : precision_score(val_labels, top1, average='macro', zero_division=0),\n        'Recall'    : recall_score(val_labels,    top1, average='macro', zero_division=0),\n        'F1'        : f1_score(val_labels,        top1, average='macro', zero_division=0),\n    }\n\neff_metrics  = retrieval_metrics(eff_val_embs,  eff_train_embs,  val_labels, train_labels, 'EfficientNetB3')\nswin_metrics = retrieval_metrics(swin_val_embs, swin_train_embs, val_labels, train_labels, 'Swin-Base')\n\nmetrics_df = pd.DataFrame([eff_metrics, swin_metrics]).set_index('Model')\nprint(metrics_df.to_string(float_format='{:.4f}'.format))\n\n# Set both variable names — submission cell uses BEST_MODEL_TYPE\nBEST = 'swin' if swin_metrics['Top-1 Acc'] >= eff_metrics['Top-1 Acc'] else 'efficientnet'\nBEST_MODEL_TYPE = BEST\nBEST_MODEL_NAME = 'Swin-Small' if BEST == 'swin' else 'EfficientNetB3'\n\nprint(f\"\\nBest model: {BEST_MODEL_NAME} ({BEST_MODEL_TYPE})\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-13T01:28:51.731680Z","iopub.execute_input":"2026-03-13T01:28:51.731891Z","iopub.status.idle":"2026-03-13T01:28:51.781263Z","shell.execute_reply.started":"2026-03-13T01:28:51.731873Z","shell.execute_reply":"2026-03-13T01:28:51.780381Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"keys  = ['Top-1 Acc', 'Top-5 Acc', 'Precision', 'Recall', 'F1']\nx     = np.arange(len(keys)); w = 0.35\n\nfig, ax = plt.subplots(figsize=(12, 5))\nax.bar(x-w/2, [eff_metrics[k]  for k in keys], w, label='EfficientNetB3', color='steelblue')\nax.bar(x+w/2, [swin_metrics[k] for k in keys], w, label='Swin-Base',      color='darkorange')\nax.set_xticks(x); ax.set_xticklabels(keys)\nax.set_ylim(0, 1.1); ax.legend(); ax.grid(axis='y', alpha=0.3)\nax.set_title('Model Comparison', fontweight='bold')\nplt.tight_layout()\nplt.savefig(CFG['OUTPUT_DIR'] / 'comparison.png', dpi=150)\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-13T01:28:51.782138Z","iopub.execute_input":"2026-03-13T01:28:51.782363Z","iopub.status.idle":"2026-03-13T01:28:52.084032Z","shell.execute_reply.started":"2026-03-13T01:28:51.782344Z","shell.execute_reply":"2026-03-13T01:28:52.083324Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def visualize_embeddings(embs, labels, name, max_cls=20, n=600):\n    unique = np.unique(labels)[:max_cls]\n    mask   = np.isin(labels, unique)\n    e, l   = embs[mask], labels[mask]\n    if len(e) > n:\n        idx = np.random.choice(len(e), n, replace=False)\n        e, l = e[idx], l[idx]\n\n    pal    = sns.color_palette('tab20', n_colors=len(unique))\n    lut    = {u: pal[i] for i, u in enumerate(unique)}\n    colors = [lut[x] for x in l]\n\n    fig, axes = plt.subplots(1, 2, figsize=(16, 6))\n    fig.suptitle(f'{name} — Embeddings', fontweight='bold')\n\n    tsne = TSNE(n_components=2, perplexity=30, random_state=SEED, n_iter=500).fit_transform(e)\n    axes[0].scatter(tsne[:,0], tsne[:,1], c=colors, s=12, alpha=0.7)\n    axes[0].set_title('t-SNE')\n\n    um = umap.UMAP(n_components=2, random_state=SEED).fit_transform(e)\n    axes[1].scatter(um[:,0], um[:,1], c=colors, s=12, alpha=0.7)\n    axes[1].set_title('UMAP')\n\n    plt.tight_layout()\n    plt.savefig(CFG['OUTPUT_DIR'] / f'emb_{name}.png', dpi=150)\n    plt.show()\n\nvisualize_embeddings(eff_val_embs,  val_labels, 'EfficientNetB3')\nvisualize_embeddings(swin_val_embs, val_labels, 'Swin')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-13T01:28:52.086132Z","iopub.execute_input":"2026-03-13T01:28:52.086385Z","iopub.status.idle":"2026-03-13T01:29:02.114770Z","shell.execute_reply.started":"2026-03-13T01:28:52.086366Z","shell.execute_reply":"2026-03-13T01:29:02.114001Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"=\" * 60)\nprint(\"RESULTS\")\nprint(\"=\" * 60)\nprint(metrics_df.to_string(float_format='{:.4f}'.format))\nprint(\"=\" * 60)\nprint(f\"Best model : {BEST}\")\nprint(f\"Submission : {CFG['OUTPUT_DIR'] / 'submission.csv'}\")\nprint(\"=\" * 60)\nfor f in sorted(CFG['OUTPUT_DIR'].glob('*')):\n    print(f\"  {f.name:<40} {f.stat().st_size/1024:>8.1f} KB\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-13T01:29:02.115760Z","iopub.execute_input":"2026-03-13T01:29:02.116055Z","iopub.status.idle":"2026-03-13T01:29:02.125502Z","shell.execute_reply.started":"2026-03-13T01:29:02.116035Z","shell.execute_reply":"2026-03-13T01:29:02.124773Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Submittion**","metadata":{}},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ── Load test pairs ──────────────────────────────────────────────────────────\ntest_df  = pd.read_csv(CFG['TEST_CSV'])\nsub_df   = pd.read_csv(CFG['SUB_CSV'])\n\nprint(f\"Test pairs        : {len(test_df)}\")\nprint(f\"Unique test images: {test_df['query_image'].nunique()}\")\nprint(test_df.head())\n\n# ── Get unique test images ───────────────────────────────────────────────────\nunique_test_imgs = sorted(\n    set(test_df['query_image']) | set(test_df['gallery_image'])\n)\nprint(f\"\\nExtracting embeddings for {len(unique_test_imgs)} unique test images...\")\n\n# ── Build a df of unique test images for extraction ──────────────────────────\ntest_img_df = pd.DataFrame({\n    'image_path': [str(CFG['TEST_DIR'] / f) for f in unique_test_imgs],\n    'filename'  : unique_test_imgs\n})\n\n# ── Extract embeddings with best model ───────────────────────────────────────\nif BEST_MODEL_TYPE == 'efficientnet':\n    test_embs = extract_eff(test_img_df)\nelif BEST_MODEL_TYPE == 'swin':\n    test_embs = extract_swin(test_img_df)\nelse:\n    # Ensemble\n    eff_te  = extract_eff(test_img_df)\n    swin_te = extract_swin(test_img_df)\n    concat  = np.concatenate([eff_te, swin_te], axis=1).astype(np.float32)\n    _, _, Vt = np.linalg.svd(concat, full_matrices=False)\n    proj    = concat @ Vt[:CFG['EMBEDDING_DIM']].T\n    norms   = np.linalg.norm(proj, axis=1, keepdims=True)\n    test_embs = proj / (norms + 1e-8)\n\n# L2 normalize all embeddings\nnorms      = np.linalg.norm(test_embs, axis=1, keepdims=True)\ntest_embs  = test_embs / (norms + 1e-8)\n\n# ── Build filename → embedding lookup ────────────────────────────────────────\nemb_lookup = {\n    fname: test_embs[i]\n    for i, fname in enumerate(unique_test_imgs)\n}\n\nprint(f\"Embedding shape per image: {test_embs[0].shape}\")\n\n# ── Compute cosine similarity for every pair ─────────────────────────────────\nprint(f\"\\nScoring {len(test_df):,} pairs...\")\n\nsimilarities = []\nfor _, row in tqdm(test_df.iterrows(), total=len(test_df), desc='Scoring pairs'):\n    q_emb = emb_lookup[row['query_image']]\n    g_emb = emb_lookup[row['gallery_image']]\n\n    # Cosine similarity — both already L2 normalized so dot product suffices\n    cos_sim = float(np.dot(q_emb, g_emb))\n\n    # Map from [-1, 1] → [0, 1]\n    score = (cos_sim + 1.0) / 2.0\n\n    similarities.append(score)\n\nsimilarities = np.array(similarities, dtype=np.float32)\n\n# ── Build submission ─────────────────────────────────────────────────────────\nsubmission = pd.DataFrame({\n    'row_id'    : test_df['row_id'].values,\n    'similarity': similarities\n})\n\n# ── Validate ─────────────────────────────────────────────────────────────────\nassert len(submission) == 137270,            f\"Wrong rows: {len(submission)}\"\nassert list(submission.columns) == ['row_id','similarity'], \"Wrong columns\"\nassert submission['row_id'].tolist() == list(range(137270)), \"Wrong row_id order\"\nassert np.isfinite(similarities).all(),      \"Contains NaN or Inf\"\nassert (similarities >= 0).all(),            f\"Min value: {similarities.min()}\"\nassert (similarities <= 1).all(),            f\"Max value: {similarities.max()}\"\n\nprint(\"\\n\" + \"=\"*50)\nprint(\"VALIDATION PASSED\")\nprint(\"=\"*50)\nprint(f\"Rows             : {len(submission):,}\")\nprint(f\"Similarity min   : {similarities.min():.6f}\")\nprint(f\"Similarity max   : {similarities.max():.6f}\")\nprint(f\"Similarity mean  : {similarities.mean():.6f}\")\nprint(f\"Similarity std   : {similarities.std():.6f}\")\n\n# ── Save ─────────────────────────────────────────────────────────────────────\nsub_path = CFG['OUTPUT_DIR'] / 'submission.csv'\nsubmission.to_csv(sub_path, index=False, encoding='utf-8')\n\nsize_mb = sub_path.stat().st_size / 1024**2\nprint(f\"\\nSaved → {sub_path}  ({size_mb:.2f} MB)\")\nsubmission.head(10)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-13T01:37:51.511731Z","iopub.execute_input":"2026-03-13T01:37:51.512556Z","iopub.status.idle":"2026-03-13T01:39:56.579711Z","shell.execute_reply.started":"2026-03-13T01:37:51.512528Z","shell.execute_reply":"2026-03-13T01:39:56.578811Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ── Extract test embeddings with EfficientNet ─────────────────────────────────\nprint(\"Extracting EfficientNet test embeddings...\")\ntest_eff_embs = []\nfor b in tqdm(range(math.ceil(len(test_img_df)/64)), desc='EfficientNet test'):\n    imgs = np.stack([load_and_preprocess(p)\n                     for p in test_img_df['image_path'].tolist()[b*64:(b+1)*64]])\n    test_eff_embs.append(eff_emb_model.predict(imgs, verbose=0))\ntest_eff_embs = np.vstack(test_eff_embs)\n\n# L2 normalize\nnorms         = np.linalg.norm(test_eff_embs, axis=1, keepdims=True)\ntest_eff_embs = test_eff_embs / (norms + 1e-8)\n\n# Build lookup\neff_lookup = {\n    fname: test_eff_embs[i]\n    for i, fname in enumerate(unique_test_imgs)\n}\n\nprint(f\"Test embeddings shape: {test_eff_embs.shape}\")\nprint(f\"Sample norm: {np.linalg.norm(test_eff_embs[0]):.6f}\")\n\n# ── Score all pairs ───────────────────────────────────────────────────────────\nprint(f\"\\nScoring {len(test_df):,} pairs...\")\neff_similarities = []\nfor _, row in tqdm(test_df.iterrows(), total=len(test_df), desc='Scoring'):\n    q_emb = eff_lookup[row['query_image']]\n    g_emb = eff_lookup[row['gallery_image']]\n    score = (float(np.dot(q_emb, g_emb)) + 1.0) / 2.0\n    eff_similarities.append(score)\n\neff_similarities = np.array(eff_similarities, dtype=np.float32)\n\n# ── Build and validate submission ─────────────────────────────────────────────\neff_submission = pd.DataFrame({\n    'row_id'    : test_df['row_id'].values,\n    'similarity': eff_similarities\n})\n\nassert len(eff_submission) == 137270\nassert (eff_similarities >= 0).all()\nassert (eff_similarities <= 1).all()\nassert np.isfinite(eff_similarities).all()\n\nprint(\"\\nEfficientNet submission statistics:\")\nprint(f\"  min  : {eff_similarities.min():.6f}\")\nprint(f\"  max  : {eff_similarities.max():.6f}\")\nprint(f\"  mean : {eff_similarities.mean():.6f}\")\nprint(f\"  std  : {eff_similarities.std():.6f}\")\n\n# ── Save ──────────────────────────────────────────────────────────────────────\neff_sub_path = CFG['OUTPUT_DIR'] / 'submission_efficientnet.csv'\neff_submission.to_csv(str(eff_sub_path), index=False, encoding='utf-8')\nprint(f\"\\nSaved → {eff_sub_path}  \"\n      f\"({eff_sub_path.stat().st_size/1024**2:.2f} MB)\")\n\neff_submission.head(10)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-13T01:42:19.027906Z","iopub.execute_input":"2026-03-13T01:42:19.028468Z","iopub.status.idle":"2026-03-13T01:44:14.864509Z","shell.execute_reply.started":"2026-03-13T01:42:19.028441Z","shell.execute_reply":"2026-03-13T01:44:14.863861Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.metrics import average_precision_score\n\ndef compute_map(val_embs, train_embs, val_labels, train_labels, model_name='Model'):\n    \"\"\"\n    Identity-balanced mAP — same metric as the competition.\n    Each jaguar identity contributes equal weight regardless of image count.\n    \"\"\"\n    # L2 normalize\n    def normalize(e):\n        n = np.linalg.norm(e, axis=1, keepdims=True)\n        return e / (n + 1e-8)\n\n    q_embs = normalize(val_embs.astype(np.float32))\n    g_embs = normalize(train_embs.astype(np.float32))\n\n    # Cosine similarity matrix (n_queries x n_gallery)\n    sim_matrix = np.matmul(q_embs, g_embs.T)\n\n    # Per-query AP\n    query_ap   = {}\n    for i in range(len(val_labels)):\n        q_label  = val_labels[i]\n        scores   = sim_matrix[i]\n        gt       = (train_labels == q_label).astype(int)\n\n        if gt.sum() == 0:\n            continue\n\n        ap = average_precision_score(gt, scores)\n        if q_label not in query_ap:\n            query_ap[q_label] = []\n        query_ap[q_label].append(ap)\n\n    # Identity-balanced mAP — mean per identity first, then mean across identities\n    per_identity_map = {identity: np.mean(aps) for identity, aps in query_ap.items()}\n    balanced_map     = np.mean(list(per_identity_map.values()))\n\n    print(f\"\\n{'='*50}\")\n    print(f\"  {model_name}\")\n    print(f\"{'='*50}\")\n    print(f\"  Identity-balanced mAP : {balanced_map:.4f}\")\n    print(f\"  Identities evaluated  : {len(per_identity_map)}\")\n    print(f\"  Queries evaluated     : {sum(len(v) for v in query_ap.values())}\")\n    print(f\"  Best  identity mAP    : {max(per_identity_map.values()):.4f}\")\n    print(f\"  Worst identity mAP    : {min(per_identity_map.values()):.4f}\")\n    print(f\"{'='*50}\")\n\n    return balanced_map, per_identity_map\n\n# ── Compute for both models ───────────────────────────────────────────────────\neff_map,  eff_per_id  = compute_map(\n    eff_val_embs,  eff_train_embs,\n    val_labels,    train_labels,\n    'EfficientNetB3'\n)\n\nswin_map, swin_per_id = compute_map(\n    swin_val_embs, swin_train_embs,\n    val_labels,    train_labels,\n    'Swin-Small'\n)\n\n# ── Ensemble score ────────────────────────────────────────────────────────────\nensemble_val_embs   = np.concatenate([eff_val_embs,   swin_val_embs],   axis=1)\nensemble_train_embs = np.concatenate([eff_train_embs, swin_train_embs], axis=1)\n\nens_map, ens_per_id = compute_map(\n    ensemble_val_embs, ensemble_train_embs,\n    val_labels,        train_labels,\n    'Ensemble (EfficientNet + Swin)'\n)\n\n# ── Summary ───────────────────────────────────────────────────────────────────\nprint(f\"\\nFinal comparison:\")\nprint(f\"  EfficientNetB3  mAP : {eff_map:.4f}\")\nprint(f\"  Swin-Small      mAP : {swin_map:.4f}\")\nprint(f\"  Ensemble        mAP : {ens_map:.4f}\")\n\nbest = max([('efficientnet', eff_map), ('swin', swin_map), ('ensemble', ens_map)],\n           key=lambda x: x[1])\nprint(f\"\\n  Best → {best[0]} ({best[1]:.4f})\")\nBEST_MODEL_TYPE = best[0]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-13T01:44:14.876588Z","iopub.execute_input":"2026-03-13T01:44:14.876833Z","iopub.status.idle":"2026-03-13T01:44:15.966994Z","shell.execute_reply.started":"2026-03-13T01:44:14.876813Z","shell.execute_reply":"2026-03-13T01:44:15.966254Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import subprocess\nresult = subprocess.run([\n    'kaggle', 'competitions', 'submit',\n    '-c', 'jaguar-re-id',\n    '-f', str(CFG['OUTPUT_DIR'] / 'submission.csv'),\n    '-m', 'Swin-Small ArcFace mAP 0.9159'\n], capture_output=True, text=True)\nprint(result.stdout)\nprint(result.stderr)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-13T01:50:48.396923Z","iopub.execute_input":"2026-03-13T01:50:48.397659Z","iopub.status.idle":"2026-03-13T01:50:50.394007Z","shell.execute_reply.started":"2026-03-13T01:50:48.397607Z","shell.execute_reply":"2026-03-13T01:50:50.393329Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}