{"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":"gpu","dataSources":[{"sourceId":128792,"databundleVersionId":15494745,"sourceType":"competition"}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =============================================================================\n# 🏆 VISTA CODEFEST'26 — KAGGLE P100 FINAL BUILD (WITH RESUME MODE)\n# =============================================================================\n# FEATURES:\n#   ✅ All micro-fixes applied\n#   ✅ Resume mode (survives kernel disconnects)\n#   ✅ Checkpoint saving to /kaggle/temp\n#   ✅ Automatic stage detection on restart\n#   ✅ Symlink with fallback\n#   ✅ RAM cache fix\n#   ✅ Class mapping safety\n# =============================================================================\n\nimport os\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\nos.environ[\"CUDA_LAUNCH_BLOCKING\"] = \"0\"\nos.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\n\nimport subprocess\nimport sys\nimport json\nimport time\nimport random\nimport gc\nimport shutil\nimport pickle\nfrom pathlib import Path\nfrom collections import defaultdict, Counter\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\n# =============================================================================\n# RESUME MODE INFRASTRUCTURE\n# =============================================================================\nCHECKPOINT_DIR = \"/kaggle/temp/vista_checkpoints\"\nos.makedirs(CHECKPOINT_DIR, exist_ok=True)\n\nclass ResumeManager:\n    \"\"\"Handles checkpointing and resume logic\"\"\"\n    \n    def __init__(self, checkpoint_dir):\n        self.dir = checkpoint_dir\n        self.state_file = f\"{checkpoint_dir}/state.json\"\n        self.state = self._load_state()\n    \n    def _load_state(self):\n        if os.path.exists(self.state_file):\n            with open(self.state_file) as f:\n                return json.load(f)\n        return {\"stage\": 0, \"yolo_done\": False, \"count_done\": False, \"cal_done\": False}\n    \n    def save_state(self):\n        with open(self.state_file, \"w\") as f:\n            json.dump(self.state, f)\n    \n    def mark_stage(self, stage_num):\n        self.state[\"stage\"] = stage_num\n        self.save_state()\n    \n    def mark_yolo_done(self):\n        self.state[\"yolo_done\"] = True\n        self.save_state()\n    \n    def mark_count_done(self):\n        self.state[\"count_done\"] = True\n        self.save_state()\n    \n    def mark_cal_done(self):\n        self.state[\"cal_done\"] = True\n        self.save_state()\n    \n    def is_yolo_done(self):\n        return self.state.get(\"yolo_done\", False) and os.path.exists(f\"{self.dir}/yolo_best.pt\")\n    \n    def is_count_done(self):\n        return self.state.get(\"count_done\", False) and os.path.exists(f\"{self.dir}/count_head.pt\")\n    \n    def is_cal_done(self):\n        return self.state.get(\"cal_done\", False) and os.path.exists(f\"{self.dir}/calibration.pkl\")\n    \n    def save_calibration(self, temperature, weight):\n        with open(f\"{self.dir}/calibration.pkl\", \"wb\") as f:\n            pickle.dump({\"temperature\": temperature, \"weight\": weight}, f)\n    \n    def load_calibration(self):\n        with open(f\"{self.dir}/calibration.pkl\", \"rb\") as f:\n            return pickle.load(f)\n    \n    def reset(self):\n        \"\"\"Force fresh start\"\"\"\n        self.state = {\"stage\": 0, \"yolo_done\": False, \"count_done\": False, \"cal_done\": False}\n        self.save_state()\n\nresume_mgr = ResumeManager(CHECKPOINT_DIR)\n\n# =============================================================================\n# RUNTIME WATCHDOG\n# =============================================================================\nKAGGLE_TIME_LIMIT = 8.0 * 60 * 60\nSTART_TIME = time.time()\n\ndef time_remaining():\n    return KAGGLE_TIME_LIMIT - (time.time() - START_TIME)\n\ndef time_ok(required_minutes=30):\n    return time_remaining() > required_minutes * 60\n\ndef elapsed():\n    return f\"{(time.time()-START_TIME)/60:.1f}m\"\n\n# =============================================================================\n# DEPENDENCIES\n# =============================================================================\ndef install(pkg):\n    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", pkg])\n\ntry:\n    import ultralytics\nexcept ImportError:\n    install(\"ultralytics\")\n\ntry:\n    import timm\nexcept ImportError:\n    install(\"timm\")\n\nimport yaml\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.cuda.amp import autocast, GradScaler\nfrom torchvision import transforms\nfrom ultralytics import YOLO\nimport timm\n\n# =============================================================================\n# GPU SETUP\n# =============================================================================\nprint(\"=\" * 70)\nprint(\"🏆 VISTA P100 FINAL BUILD (RESUME MODE)\")\nprint(\"=\" * 70)\n\nif not torch.cuda.is_available():\n    raise RuntimeError(\"❌ GPU REQUIRED\")\n\nDEVICE = torch.device(\"cuda:0\")\ntorch.backends.cudnn.benchmark = True\ntorch.backends.cuda.matmul.allow_tf32 = True\n\n# CUDA warmup\n_ = torch.randn(1, 3, 640, 640, device=DEVICE)\n_ = torch.randn(1, 3, 224, 224, device=DEVICE)\ntorch.cuda.synchronize()\n\nprint(f\"✅ GPU: {torch.cuda.get_device_name(0)}\")\nprint(f\"✅ VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.2f} GB\")\n\n# Check resume state\nif resume_mgr.is_yolo_done():\n    print(\"📌 RESUME: YOLO checkpoint found\")\nif resume_mgr.is_count_done():\n    print(\"📌 RESUME: Count head checkpoint found\")\nif resume_mgr.is_cal_done():\n    print(\"📌 RESUME: Calibration checkpoint found\")\n\n# =============================================================================\n# CONFIG\n# =============================================================================\nclass Cfg:\n    ROOT = \"/kaggle/input/vista26\"\n    BASE = f\"{ROOT}/Vistas Dataset Public/Vistas Dataset Public\"\n    WORK = \"/kaggle/working\"\n    \n    TRAIN_DIR = f\"{BASE}/train\"\n    TEST_DIR = f\"{BASE}/test\"\n    VAL_DIR = f\"{BASE}/validation\"\n    \n    TRAIN_JSON = f\"{BASE}/instances_train.json\"\n    TEST_JSON = f\"{BASE}/instances_test.json\"\n    VAL_JSON = f\"{ROOT}/instances_val.json\"\n    CATS_JSON = f\"{BASE}/Categories.json\"\n    \n    # YOLO — KAGGLE SAFE\n    YOLO_MODEL = \"yolov8s.pt\"\n    YOLO_EPOCHS = 18\n    YOLO_IMGSZ = 640\n    YOLO_BATCH = 16\n    YOLO_WORKERS = 0\n    YOLO_PATIENCE = 6\n    \n    # Count Head — KAGGLE SAFE\n    COUNT_BACKBONE = \"efficientnet_b0\"\n    COUNT_IMGSZ = 224\n    COUNT_EPOCHS = 16\n    COUNT_BATCH = 32\n    COUNT_LR = 1e-3\n    MAX_COUNT = 25\n    \n    # Data\n    TRAIN_SAMPLE_SIZE = 7000\n    \n    # Inference\n    CONF_FLOOR = 0.003\n    IOU = 0.45\n    MAX_DET = 40\n    \n    # Defaults\n    CONF_MASS_WEIGHT = 0.35\n    TEMPERATURE = 1.2\n    \n    SEED = 42\n\nC = Cfg()\nrandom.seed(C.SEED)\nnp.random.seed(C.SEED)\ntorch.manual_seed(C.SEED)\ntorch.cuda.manual_seed_all(C.SEED)\n\ndef cleanup():\n    gc.collect()\n    torch.cuda.empty_cache()\n    torch.cuda.synchronize()\n\n# =============================================================================\n# STAGE 0: LOAD CATEGORIES\n# =============================================================================\nprint(f\"\\n[{elapsed()}] STAGE 0 — Loading Categories\")\n\nwith open(C.CATS_JSON) as f:\n    cats_data = json.load(f)[\"categories\"]\n\ncats_data.sort(key=lambda c: int(c[\"id\"]))\nORIGINAL_IDS = [int(c[\"id\"]) for c in cats_data]\nNUM_CLASSES = len(cats_data)\n\norig_to_yolo = {orig_id: idx for idx, orig_id in enumerate(ORIGINAL_IDS)}\nyolo_to_orig = {idx: orig_id for idx, orig_id in enumerate(ORIGINAL_IDS)}\nnames_dict = {idx: cats_data[idx][\"name\"] for idx in range(NUM_CLASSES)}\nVALID_IDS = set(ORIGINAL_IDS)\n\nprint(f\"✅ {NUM_CLASSES} categories\")\n\nLEVEL_TO_IDX = {\"easy\": 0, \"medium\": 1, \"hard\": 2}\n\n# =============================================================================\n# STAGE 1: LOAD DATA\n# =============================================================================\nprint(f\"\\n[{elapsed()}] STAGE 1 — Loading Data\")\n\ndef load_json(path):\n    with open(path) as f:\n        return json.load(f)\n\ntrain_data = load_json(C.TRAIN_JSON)\ntest_data = load_json(C.TEST_JSON)\n\ntrain_images = {}\nfor img in train_data[\"images\"]:\n    train_images[int(img[\"id\"])] = {\n        \"file_name\": img[\"file_name\"],\n        \"width\": int(img[\"width\"]),\n        \"height\": int(img[\"height\"]),\n        \"level\": \"easy\"\n    }\n\ntrain_anns = defaultdict(list)\nfor ann in train_data[\"annotations\"]:\n    cat_id = int(ann[\"category_id\"])\n    if cat_id in VALID_IDS:\n        train_anns[int(ann[\"image_id\"])].append({\n            \"category_id\": cat_id,\n            \"bbox\": ann[\"bbox\"]\n        })\n\ntest_images = {}\nfor img in test_data[\"images\"]:\n    test_images[int(img[\"id\"])] = {\n        \"file_name\": img[\"file_name\"],\n        \"width\": int(img[\"width\"]),\n        \"height\": int(img[\"height\"]),\n        \"level\": img.get(\"level\", \"medium\")\n    }\n\ntest_anns = defaultdict(list)\nfor ann in test_data[\"annotations\"]:\n    cat_id = int(ann[\"category_id\"])\n    if cat_id in VALID_IDS:\n        test_anns[int(ann[\"image_id\"])].append({\n            \"category_id\": cat_id,\n            \"bbox\": ann[\"bbox\"]\n        })\n\ntrain_ids = [i for i in train_images if train_anns[i]]\ntest_ids = [i for i in test_images if test_anns[i]]\n\n# STRICT 4-WAY SPLIT\nrandom.shuffle(test_ids)\n\nby_level = defaultdict(list)\nfor iid in test_ids:\n    by_level[test_images[iid][\"level\"]].append(iid)\n\nyolo_train_ids = []\ncount_train_ids = []\ntemp_cal_ids = []\nselection_cal_ids = []\n\nfor level in [\"easy\", \"medium\", \"hard\"]:\n    ids = by_level[level]\n    random.shuffle(ids)\n    n = len(ids)\n    s1, s2, s3 = int(n * 0.65), int(n * 0.85), int(n * 0.93)\n    yolo_train_ids.extend(ids[:s1])\n    count_train_ids.extend(ids[s1:s2])\n    temp_cal_ids.extend(ids[s2:s3])\n    selection_cal_ids.extend(ids[s3:])\n\nprint(f\"✅ YOLO: {len(yolo_train_ids)} | Count: {len(count_train_ids)} | TempCal: {len(temp_cal_ids)} | SelCal: {len(selection_cal_ids)}\")\n\n# =============================================================================\n# MODEL DEFINITIONS\n# =============================================================================\nclass CountDistributionHead(nn.Module):\n    def __init__(self, backbone_name=\"efficientnet_b0\", max_count=25):\n        super().__init__()\n        self.max_count = max_count\n        self.num_classes = max_count + 1\n        \n        self.backbone = timm.create_model(backbone_name, pretrained=True, num_classes=0)\n        feat_dim = self.backbone.num_features\n        \n        self.level_embed = nn.Embedding(3, 64)\n        \n        self.head = nn.Sequential(\n            nn.Linear(feat_dim + 64, 512),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(512, 256),\n            nn.ReLU(),\n            nn.Dropout(0.2),\n            nn.Linear(256, self.num_classes)\n        )\n    \n    def forward(self, x, level_idx):\n        features = self.backbone(x)\n        level_emb = self.level_embed(level_idx)\n        combined = torch.cat([features, level_emb], dim=1)\n        return self.head(combined)\n\n\nclass LabelSmoothingLoss(nn.Module):\n    def __init__(self, num_classes, smoothing=0.1):\n        super().__init__()\n        self.num_classes = num_classes\n        self.smoothing = smoothing\n        self.confidence = 1.0 - smoothing\n    \n    def forward(self, pred, target):\n        pred = F.log_softmax(pred, dim=1)\n        \n        with torch.no_grad():\n            true_dist = torch.zeros_like(pred)\n            true_dist.scatter_(1, target.unsqueeze(1), self.confidence)\n            \n            smooth_mass = self.smoothing / 2\n            for i, t in enumerate(target):\n                if t > 0:\n                    true_dist[i, t-1] += smooth_mass\n                else:\n                    true_dist[i, t] += smooth_mass\n                if t < self.num_classes - 1:\n                    true_dist[i, t+1] += smooth_mass\n                else:\n                    true_dist[i, t] += smooth_mass\n        \n        return torch.mean(torch.sum(-true_dist * pred, dim=1))\n\n\nclass TemperatureScaling:\n    def __init__(self):\n        self.temperature = 1.0\n    \n    def fit(self, logits_list, labels_list):\n        logits = torch.tensor(np.array(logits_list), dtype=torch.float32)\n        labels = torch.tensor(np.array(labels_list), dtype=torch.long)\n        \n        best_nll, best_temp = float('inf'), 1.0\n        \n        for temp in np.arange(0.5, 2.5, 0.05):\n            nll = F.nll_loss(F.log_softmax(logits / temp, dim=1), labels).item()\n            if nll < best_nll:\n                best_nll, best_temp = nll, temp\n        \n        for temp in np.arange(max(0.3, best_temp - 0.1), best_temp + 0.1, 0.01):\n            nll = F.nll_loss(F.log_softmax(logits / temp, dim=1), labels).item()\n            if nll < best_nll:\n                best_nll, best_temp = nll, temp\n        \n        self.temperature = best_temp\n        acc = (F.softmax(logits / self.temperature, dim=1).argmax(dim=1) == labels).float().mean().item()\n        print(f\"  Temperature: {self.temperature:.3f} | Accuracy: {acc:.3f}\")\n        return self.temperature\n    \n    def calibrate(self, logits):\n        return F.softmax(torch.tensor(logits) / self.temperature, dim=-1).numpy()\n\n# =============================================================================\n# STAGE 2: PREPARE YOLO DATASET\n# =============================================================================\nprint(f\"\\n[{elapsed()}] STAGE 2 — Preparing YOLO Dataset\")\n\nYOLO_DIR = f\"{C.WORK}/yolo_data\"\n\ndef write_yolo_safe(iid, split, img_info, anns, img_dir, prefix=\"\"):\n    \"\"\"Symlink with fallback to copy\"\"\"\n    src = f\"{img_dir}/{img_info['file_name']}\"\n    if not os.path.exists(src):\n        return False\n    \n    stem = f\"{prefix}{iid}\"\n    img_dst = f\"{YOLO_DIR}/images/{split}/{stem}.jpg\"\n    lbl_dst = f\"{YOLO_DIR}/labels/{split}/{stem}.txt\"\n    \n    # SYMLINK WITH FALLBACK (FIX #1)\n    if not os.path.exists(img_dst):\n        try:\n            os.symlink(src, img_dst)\n        except (OSError, NotImplementedError):\n            shutil.copy2(src, img_dst)\n    \n    W, H = img_info[\"width\"], img_info[\"height\"]\n    lines = []\n    for ann in anns:\n        cls = orig_to_yolo[ann[\"category_id\"]]\n        x, y, w, h = ann[\"bbox\"]\n        cx = np.clip((x + w/2) / W, 0.001, 0.999)\n        cy = np.clip((y + h/2) / H, 0.001, 0.999)\n        nw = np.clip(w / W, 0.001, 0.999)\n        nh = np.clip(h / H, 0.001, 0.999)\n        lines.append(f\"{cls} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}\")\n    \n    with open(lbl_dst, \"w\") as f:\n        f.write(\"\\n\".join(lines))\n    return True\n\n# Only rebuild if needed\nif not os.path.exists(f\"{YOLO_DIR}/images/train\") or len(os.listdir(f\"{YOLO_DIR}/images/train\")) < 100:\n    shutil.rmtree(YOLO_DIR, ignore_errors=True)\n    for sub in [\"images/train\", \"images/val\", \"labels/train\", \"labels/val\"]:\n        os.makedirs(f\"{YOLO_DIR}/{sub}\", exist_ok=True)\n    \n    random.shuffle(train_ids)\n    train_sample = train_ids[:C.TRAIN_SAMPLE_SIZE]\n    \n    for iid in train_sample[:int(len(train_sample)*0.9)]:\n        write_yolo_safe(iid, \"train\", train_images[iid], train_anns[iid], C.TRAIN_DIR, \"s_\")\n    for iid in train_sample[int(len(train_sample)*0.9):]:\n        write_yolo_safe(iid, \"val\", train_images[iid], train_anns[iid], C.TRAIN_DIR, \"s_\")\n    \n    for iid in yolo_train_ids:\n        write_yolo_safe(iid, \"train\", test_images[iid], test_anns[iid], C.TEST_DIR, \"m_\")\n    for iid in temp_cal_ids + selection_cal_ids:\n        write_yolo_safe(iid, \"val\", test_images[iid], test_anns[iid], C.TEST_DIR, \"m_\")\n    \n    print(f\"✅ Dataset built\")\nelse:\n    print(f\"✅ Dataset exists, reusing\")\n\nds_yaml = {\"path\": C.WORK, \"train\": \"yolo_data/images/train\", \"val\": \"yolo_data/images/val\", \"nc\": NUM_CLASSES, \"names\": names_dict}\nwith open(f\"{C.WORK}/dataset.yaml\", \"w\") as f:\n    yaml.dump(ds_yaml, f)\n\n# =============================================================================\n# STAGE 3: TRAIN YOLO (WITH RESUME)\n# =============================================================================\nprint(f\"\\n[{elapsed()}] STAGE 3 — YOLO Training\")\n\nBEST_YOLO = f\"{CHECKPOINT_DIR}/yolo_best.pt\"\n\nif resume_mgr.is_yolo_done():\n    print(f\"📌 RESUME: Loading YOLO from checkpoint\")\n    yolo_model = YOLO(BEST_YOLO)\nelse:\n    yolo_model = YOLO(C.YOLO_MODEL)\n    yolo_model.train(\n        data=f\"{C.WORK}/dataset.yaml\",\n        epochs=C.YOLO_EPOCHS,\n        imgsz=C.YOLO_IMGSZ,\n        batch=C.YOLO_BATCH,\n        device=0,\n        workers=C.YOLO_WORKERS,\n        project=C.WORK,\n        name=\"yolo_run\",\n        exist_ok=True,\n        patience=C.YOLO_PATIENCE,\n        amp=True,\n        cache=\"ram\",  # FIX #2: Explicit RAM cache\n        mosaic=0.6,\n        mixup=0.1,\n        close_mosaic=4,\n        verbose=False,\n    )\n    \n    # Find best weights\n    yolo_best = f\"{C.WORK}/yolo_run/weights/best.pt\"\n    if not os.path.exists(yolo_best):\n        yolo_best = f\"{C.WORK}/yolo_run/weights/last.pt\"\n    \n    # Save to checkpoint dir\n    shutil.copy2(yolo_best, BEST_YOLO)\n    resume_mgr.mark_yolo_done()\n    yolo_model = YOLO(BEST_YOLO)\n\nprint(f\"✅ YOLO ready\")\ncleanup()\n\n# =============================================================================\n# STAGE 4: TRAIN COUNT HEAD (WITH RESUME + AMP)\n# =============================================================================\nprint(f\"\\n[{elapsed()}] STAGE 4 — Count Head Training\")\n\ncount_transform = transforms.Compose([\n    transforms.ToPILImage(),\n    transforms.Resize((C.COUNT_IMGSZ, C.COUNT_IMGSZ)),\n    transforms.ToTensor(),\n    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n])\n\ncount_transform_aug = transforms.Compose([\n    transforms.ToPILImage(),\n    transforms.Resize((C.COUNT_IMGSZ + 32, C.COUNT_IMGSZ + 32)),\n    transforms.RandomCrop(C.COUNT_IMGSZ),\n    transforms.RandomHorizontalFlip(),\n    transforms.ColorJitter(brightness=0.2, contrast=0.2),\n    transforms.ToTensor(),\n    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n])\n\nclass CountDataset(Dataset):\n    def __init__(self, ids, images_dict, anns_dict, img_dir, transform, max_count):\n        self.samples = []\n        for iid in ids:\n            info = images_dict[iid]\n            path = f\"{img_dir}/{info['file_name']}\"\n            if os.path.exists(path):\n                self.samples.append({\n                    \"path\": path,\n                    \"count\": min(len(anns_dict[iid]), max_count),\n                    \"level\": info.get(\"level\", \"medium\")\n                })\n        self.transform = transform\n    \n    def __len__(self):\n        return len(self.samples)\n    \n    def __getitem__(self, idx):\n        s = self.samples[idx]\n        img = cv2.imread(s[\"path\"])\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        img = self.transform(img)\n        return img, torch.tensor(LEVEL_TO_IDX.get(s[\"level\"], 1)), torch.tensor(s[\"count\"], dtype=torch.long)\n\ncount_model = CountDistributionHead(C.COUNT_BACKBONE, C.MAX_COUNT).to(DEVICE)\nCOUNT_CKPT = f\"{CHECKPOINT_DIR}/count_head.pt\"\n\nif resume_mgr.is_count_done():\n    print(f\"📌 RESUME: Loading Count Head from checkpoint\")\n    state = torch.load(COUNT_CKPT, map_location=DEVICE)\n    count_model.load_state_dict(state)\n    best_acc = 0.5  # Placeholder\nelse:\n    split_idx = int(len(count_train_ids) * 0.85)\n    count_train_ds = CountDataset(count_train_ids[:split_idx], test_images, test_anns, C.TEST_DIR, count_transform_aug, C.MAX_COUNT)\n    count_val_ds = CountDataset(count_train_ids[split_idx:], test_images, test_anns, C.TEST_DIR, count_transform, C.MAX_COUNT)\n    \n    count_train_loader = DataLoader(count_train_ds, batch_size=C.COUNT_BATCH, shuffle=True, num_workers=0, pin_memory=True)\n    count_val_loader = DataLoader(count_val_ds, batch_size=C.COUNT_BATCH, shuffle=False, num_workers=0, pin_memory=True)\n    \n    optimizer = torch.optim.AdamW(count_model.parameters(), lr=C.COUNT_LR, weight_decay=1e-4)\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, C.COUNT_EPOCHS)\n    criterion = LabelSmoothingLoss(C.MAX_COUNT + 1, smoothing=0.1)\n    scaler = GradScaler()\n    \n    best_acc = 0.0\n    best_state = None\n    \n    for epoch in range(C.COUNT_EPOCHS):\n        count_model.train()\n        for imgs, levels, counts in count_train_loader:\n            imgs, levels, counts = imgs.to(DEVICE), levels.to(DEVICE), counts.to(DEVICE)\n            \n            optimizer.zero_grad()\n            with autocast():\n                logits = count_model(imgs, levels)\n                loss = criterion(logits, counts)\n            \n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n        \n        scheduler.step()\n        \n        count_model.eval()\n        correct, total = 0, 0\n        with torch.no_grad():\n            for imgs, levels, counts in count_val_loader:\n                imgs, levels, counts = imgs.to(DEVICE), levels.to(DEVICE), counts.to(DEVICE)\n                with autocast():\n                    logits = count_model(imgs, levels)\n                preds = logits.argmax(dim=1)\n                correct += (preds == counts).sum().item()\n                total += counts.size(0)\n        \n        acc = correct / total\n        if acc > best_acc:\n            best_acc = acc\n            best_state = {k: v.cpu().clone() for k, v in count_model.state_dict().items()}\n        \n        if (epoch + 1) % 4 == 0:\n            print(f\"  Epoch {epoch+1}: Acc={acc:.3f}\")\n    \n    count_model.load_state_dict({k: v.to(DEVICE) for k, v in best_state.items()})\n    torch.save(best_state, COUNT_CKPT)\n    resume_mgr.mark_count_done()\n\nprint(f\"✅ Count Head Accuracy: {best_acc:.3f}\")\ncleanup()\n\n# =============================================================================\n# STAGE 5: CALIBRATION (WITH RESUME)\n# =============================================================================\nprint(f\"\\n[{elapsed()}] STAGE 5 — Calibration\")\n\ntemp_scaler = TemperatureScaling()\n\nif resume_mgr.is_cal_done():\n    print(f\"📌 RESUME: Loading calibration from checkpoint\")\n    cal_data = resume_mgr.load_calibration()\n    C.TEMPERATURE = cal_data[\"temperature\"]\n    C.CONF_MASS_WEIGHT = cal_data[\"weight\"]\n    temp_scaler.temperature = C.TEMPERATURE\n    print(f\"  Temperature: {C.TEMPERATURE:.3f} | Weight: {C.CONF_MASS_WEIGHT:.2f}\")\nelse:\n    if time_ok(required_minutes=60):\n        print(\"  Running temperature calibration...\")\n        count_model.eval()\n        \n        temp_logits, temp_labels = [], []\n        for iid in temp_cal_ids:\n            info = test_images[iid]\n            path = f\"{C.TEST_DIR}/{info['file_name']}\"\n            if not os.path.exists(path):\n                continue\n            \n            img = cv2.imread(path)\n            if img is None:\n                continue\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n            img_tensor = count_transform(img).unsqueeze(0).to(DEVICE)\n            level_tensor = torch.tensor([LEVEL_TO_IDX.get(info[\"level\"], 1)]).to(DEVICE)\n            \n            with torch.no_grad(), autocast():\n                logits = count_model(img_tensor, level_tensor)\n            temp_logits.append(logits[0].float().cpu().numpy())\n            temp_labels.append(min(len(test_anns[iid]), C.MAX_COUNT))\n        \n        if len(temp_logits) > 50:\n            C.TEMPERATURE = temp_scaler.fit(temp_logits, temp_labels)\n    else:\n        print(\"  ⏱️ Time low — using defaults\")\n    \n    # Selection weight calibration\n    if time_ok(required_minutes=45):\n        print(\"  Running selection weight calibration...\")\n        \n        def compute_conf_mass(dets, n):\n            if n == 0 or not dets:\n                return 0.0\n            sorted_d = sorted(dets, key=lambda x: x[0], reverse=True)\n            return sum(d[0] for d in sorted_d[:n]) / n\n        \n        def select_count(probs, dets, weight, max_det):\n            if not dets:\n                return 0\n            best_n, best_score = 0, 0.0\n            sorted_d = sorted(dets, key=lambda x: x[0], reverse=True)\n            max_n = min(len(sorted_d), max_det, len(probs) - 1)\n            \n            for n in range(max_n + 1):\n                p = probs[n] if n < len(probs) else 0.0\n                cm = 1.0 if n == 0 else compute_conf_mass(sorted_d, n)\n                score = (p ** (1 - weight)) * (cm ** weight)\n                if n > 0 and n < len(sorted_d):\n                    gap = sorted_d[n-1][0] - sorted_d[n][0]\n                    if gap > 0.1:\n                        score *= 1.02\n                if score > best_score:\n                    best_score, best_n = score, n\n            return best_n\n        \n        cal_data = []\n        for iid in selection_cal_ids[:200]:\n            info = test_images[iid]\n            path = f\"{C.TEST_DIR}/{info['file_name']}\"\n            if not os.path.exists(path):\n                continue\n            \n            results = yolo_model.predict(path, conf=C.CONF_FLOOR, iou=C.IOU, max_det=C.MAX_DET, device=0, verbose=False)\n            res = results[0]\n            dets = []\n            if res.boxes is not None and len(res.boxes) > 0:\n                confs = res.boxes.conf.cpu().tolist()\n                classes = res.boxes.cls.cpu().int().tolist()\n                dets = list(zip(confs, classes))\n            \n            img = cv2.imread(path)\n            if img is None:\n                continue\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n            img_tensor = count_transform(img).unsqueeze(0).to(DEVICE)\n            level_tensor = torch.tensor([LEVEL_TO_IDX.get(info[\"level\"], 1)]).to(DEVICE)\n            \n            with torch.no_grad(), autocast():\n                logits = count_model(img_tensor, level_tensor)\n            probs = temp_scaler.calibrate(logits[0].float().cpu().numpy())\n            \n            gt = min(len(test_anns[iid]), C.MAX_COUNT)\n            cal_data.append({\"probs\": probs, \"dets\": dets, \"gt\": gt})\n        \n        if len(cal_data) > 30:\n            best_w, best_acc_cal = 0.35, 0.0\n            for w in np.arange(0.1, 0.6, 0.02):\n                correct = sum(1 for s in cal_data if select_count(s[\"probs\"], s[\"dets\"], w, C.MAX_COUNT) == s[\"gt\"])\n                acc_cal = correct / len(cal_data)\n                if acc_cal > best_acc_cal:\n                    best_acc_cal, best_w = acc_cal, w\n            C.CONF_MASS_WEIGHT = best_w\n            print(f\"  Optimal weight: {C.CONF_MASS_WEIGHT:.2f} | Exact: {best_acc_cal:.3f}\")\n    \n    # Save calibration\n    resume_mgr.save_calibration(C.TEMPERATURE, C.CONF_MASS_WEIGHT)\n    resume_mgr.mark_cal_done()\n\ncleanup()\n\n# =============================================================================\n# STAGE 6: FINAL INFERENCE\n# =============================================================================\nprint(f\"\\n[{elapsed()}] STAGE 6 — Final Inference\")\nprint(f\"  Time remaining: {time_remaining()/60:.0f} minutes\")\n\nwith open(C.VAL_JSON) as f:\n    val_meta = json.load(f)[\"images\"]\n\nsubmission_data = {\n    int(img[\"id\"]): {\"path\": f\"{C.VAL_DIR}/{img['file_name']}\", \"level\": img.get(\"level\", \"medium\")}\n    for img in val_meta\n}\nsubmission_ids = sorted(submission_data.keys())\n\nprint(f\"  Processing {len(submission_ids)} images...\")\n\ndef compute_conf_mass(dets, n):\n    if n == 0 or not dets:\n        return 0.0\n    sorted_d = sorted(dets, key=lambda x: x[0], reverse=True)\n    return sum(d[0] for d in sorted_d[:n]) / n\n\ndef select_count(probs, dets, weight, max_det):\n    if not dets:\n        return 0\n    best_n, best_score = 0, 0.0\n    sorted_d = sorted(dets, key=lambda x: x[0], reverse=True)\n    max_n = min(len(sorted_d), max_det, len(probs) - 1)\n    \n    for n in range(max_n + 1):\n        p = probs[n] if n < len(probs) else 0.0\n        cm = 1.0 if n == 0 else compute_conf_mass(sorted_d, n)\n        score = (p ** (1 - weight)) * (cm ** weight)\n        if n > 0 and n < len(sorted_d):\n            gap = sorted_d[n-1][0] - sorted_d[n][0]\n            if gap > 0.1:\n                score *= 1.02\n        if score > best_score:\n            best_score, best_n = score, n\n    return best_n\n\nfinal_results = {}\ncount_model.eval()\n\nfor idx, iid in enumerate(submission_ids):\n    info = submission_data[iid]\n    path = info[\"path\"]\n    \n    if not os.path.exists(path):\n        final_results[iid] = []\n        continue\n    \n    img = cv2.imread(path)\n    if img is None:\n        final_results[iid] = []\n        continue\n    \n    # Count distribution\n    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    img_tensor = count_transform(img_rgb).unsqueeze(0).to(DEVICE)\n    level_tensor = torch.tensor([LEVEL_TO_IDX.get(info[\"level\"], 1)]).to(DEVICE)\n    \n    with torch.no_grad(), autocast():\n        logits = count_model(img_tensor, level_tensor)\n    probs = temp_scaler.calibrate(logits[0].float().cpu().numpy())\n    \n    # YOLO detections\n    results = yolo_model.predict(path, conf=C.CONF_FLOOR, iou=C.IOU, max_det=C.MAX_DET, device=0, verbose=False)\n    res = results[0]\n    \n    if res.boxes is None or len(res.boxes) == 0:\n        final_results[iid] = []\n        continue\n    \n    confs = res.boxes.conf.cpu().tolist()\n    classes = res.boxes.cls.cpu().int().tolist()\n    dets = list(zip(confs, classes))\n    \n    # MAP selection\n    N = select_count(probs, dets, C.CONF_MASS_WEIGHT, C.MAX_COUNT)\n    \n    if N == 0:\n        final_results[iid] = []\n    else:\n        sorted_d = sorted(dets, key=lambda x: x[0], reverse=True)[:N]\n        # FIX #3: Safe class mapping\n        cats = sorted([\n            yolo_to_orig[int(d[1])]\n            for d in sorted_d\n            if int(d[1]) in yolo_to_orig and yolo_to_orig[int(d[1])] in VALID_IDS\n        ])\n        final_results[iid] = cats\n    \n    if (idx + 1) % 500 == 0:\n        print(f\"    {idx+1}/{len(submission_ids)} done\")\n\nfor iid in submission_ids:\n    if iid not in final_results:\n        final_results[iid] = []\n\n# =============================================================================\n# STAGE 7: CREATE SUBMISSION\n# =============================================================================\nprint(f\"\\n[{elapsed()}] STAGE 7 — Creating Submission\")\n\nrows = [{\"image_id\": iid, \"categories\": json.dumps(sorted(final_results[iid]))} for iid in submission_ids]\ndf = pd.DataFrame(rows).sort_values(\"image_id\").reset_index(drop=True)\n\n# Validation\nassert df[\"image_id\"].is_unique, \"Duplicate IDs!\"\nassert len(df) == len(submission_ids), \"Missing images!\"\nfor _, row in df.iterrows():\n    cats = json.loads(row[\"categories\"])\n    assert cats == sorted(cats), f\"Unsorted: {row['image_id']}\"\n    for c in cats:\n        assert c in VALID_IDS, f\"Invalid category: {c}\"\n\nOUT = f\"{C.WORK}/submission.csv\"\ndf.to_csv(OUT, index=False)\n\ntry:\n    shutil.copy2(OUT, f\"{CHECKPOINT_DIR}/submission_backup.csv\")\nexcept:\n    pass\n\ntotal_objs = sum(len(json.loads(r)) for r in df[\"categories\"])\nempty = sum(1 for r in df[\"categories\"] if r == \"[]\")\n\nprint(f\"\"\"\n{'='*70}\n🏆 VISTA P100 FINAL BUILD — COMPLETE\n{'='*70}\n📁 File: {OUT}\n📊 Images: {len(df)} | Objects: {total_objs} | Empty: {empty}\n\n🔬 CALIBRATION:\n   Temperature: {C.TEMPERATURE:.3f}\n   Weight: {C.CONF_MASS_WEIGHT:.2f}\n\n⏱️ Total: {elapsed()} | Remaining: {time_remaining()/60:.0f}m\n{'='*70}\n✅ Ready for submission!\n\n📌 RESUME INFO:\n   Checkpoints saved to: {CHECKPOINT_DIR}\n   If kernel disconnects, just re-run this cell.\n\"\"\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-31T20:10:02.573777Z","iopub.execute_input":"2026-01-31T20:10:02.574075Z","iopub.status.idle":"2026-01-31T23:00:25.866144Z","shell.execute_reply.started":"2026-01-31T20:10:02.574049Z","shell.execute_reply":"2026-01-31T23:00:25.865560Z"}},"outputs":[],"execution_count":null}]}