{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":108394,"databundleVersionId":13172641,"sourceType":"competition"}],"dockerImageVersionId":31090,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# H690 — Fast pipeline (45-minute target)\n# Author: ChatGPT — production-ready Kaggle notebook\n# Strategy: Precompute two augmented ResNet features per image + radial profile -> train small head quickly\n\n# -------------------------\n# 0) Imports & settings\n# -------------------------\nimport os, math, random, time\nfrom pathlib import Path\nfrom typing import List, Tuple\nfrom collections import defaultdict\n\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image, ImageDraw\nimport cv2\nfrom tqdm.auto import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms, models\n\nfrom sklearn.neighbors import NearestNeighbors\n\n# reproducibility\nSEED = 42\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed_all(SEED)\n\n# -------------------------\n# 1) Config (changeable)\n# -------------------------\nCACHE_ROOT = Path('/kaggle/working/h690_cache')   # cache for precomputed data\nCACHE_ROOT.mkdir(parents=True, exist_ok=True)\n\nPRECOMP_IMG_SIZE = 224        # smaller for fast feature extraction\nBACKBONE = 'resnet50'\nBACKBONE_FEATURE_DIM = 2048   # ResNet50 feature dim after global pool\nPROJ_DIM = 128                # projection dimension for contrastive loss\nPROFILE_ANGLES = 256\nBATCH_SIZE_PRECOMP = 256      # batch size when extracting backbone features (GPU heavy)\nBATCH_SIZE_TRAIN = 512        # batch size when training head (operates on precomputed arrays → small GPU mem)\nEPOCHS_HEAD = 3               # keep small; adjust if you have more time\nLR_HEAD = 1e-3\nTEMPERATURE = 0.5\nTOP_K = 8\nSIMILARITY_THRESHOLD = 0.88\n\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nNUM_WORKERS = max(0, (os.cpu_count() or 2) - 2)\n\nprint(\"Device:\", DEVICE, \"Num workers:\", NUM_WORKERS)\n\n# -------------------------\n# 2) Auto-detect dataset root\n# -------------------------\ndef find_data_root():\n    base = Path('/kaggle/input')\n    if base.exists():\n        # try to find h690-like folder\n        for child in base.iterdir():\n            if 'h690' in child.name.lower() or 'daxin' in child.name.lower() or 'gdppc' in child.name.lower():\n                return child\n        # fallback: if only one dataset mounted, return it\n        entries = list(base.iterdir())\n        if len(entries) == 1:\n            return entries[0]\n    # other fallbacks\n    for p in [Path('/input'), Path('/mnt/data'), Path.cwd()]:\n        if p.exists():\n            for child in p.iterdir():\n                if 'h690' in child.name.lower() or 'daxin' in child.name.lower():\n                    return child\n    return None\n\nDATA_ROOT = find_data_root()\nif DATA_ROOT is None:\n    raise RuntimeError(\"Could not auto-detect dataset under /kaggle/input. Attach the dataset and rerun.\")\nprint(\"Data root:\", DATA_ROOT)\n\n# -------------------------\n# 3) Build dataframe of images (robust)\n# -------------------------\ndef list_image_files(root):\n    exts = ('.png','.jpg','.jpeg','.tif','.tiff')\n    files = []\n    for ext in exts:\n        files.extend(sorted([str(p) for p in Path(root).rglob(f'*{ext}')]))\n    return files\n\n# try to use metadata CSV if available\ncsvs = list(DATA_ROOT.glob('*.csv')) + list(DATA_ROOT.rglob('*.csv'))\nmeta_csv = None\nfor c in csvs:\n    ln = c.name.lower()\n    if 'meta' in ln or 'metadata' in ln or 'info' in ln or 'gdppc' in ln:\n        meta_csv = c; break\nif meta_csv is None and csvs:\n    meta_csv = csvs[0]\n\nif meta_csv:\n    df = pd.read_csv(meta_csv)\n    # try different heuristics to find filenames\n    if 'image_path' in df.columns:\n        df['image_path'] = df['image_path'].astype(str)\n    elif 'filename' in df.columns:\n        df['image_path'] = df['filename'].apply(lambda x: str(DATA_ROOT / x))\n    elif 'id' in df.columns:\n        def guess(x):\n            for ext in ('.png','.jpg','.jpeg'):\n                p = DATA_ROOT / f\"{x}{ext}\"\n                if p.exists(): return str(p)\n            return str(DATA_ROOT / f\"{x}.png\")\n        df['image_path'] = df['id'].apply(guess)\n    else:\n        imgs = list_image_files(DATA_ROOT)\n        df = pd.DataFrame({'id':[Path(p).stem for p in imgs], 'image_path':imgs})\nelse:\n    imgs = list_image_files(DATA_ROOT)\n    df = pd.DataFrame({'id':[Path(p).stem for p in imgs], 'image_path':imgs})\n\n# ensure these files exist\ndf['exists'] = df['image_path'].apply(lambda x: Path(x).exists())\ndf = df[df['exists']].reset_index(drop=True)\nprint(\"Found images:\", len(df))\nif len(df) == 0:\n    raise RuntimeError(\"No images found. Check DATA_ROOT or how files are organized.\")\n\n# -------------------------\n# 4) Precompute: two augmented backbone features + radial profiles\n#    - saves: feat_a.npy, feat_b.npy, profile.npy, ids.csv\n# -------------------------\nFEAT_A_FILE = CACHE_ROOT / 'feat_a.npy'\nFEAT_B_FILE = CACHE_ROOT / 'feat_b.npy'\nPROFILE_FILE = CACHE_ROOT / 'profiles.npy'\nIDS_FILE = CACHE_ROOT / 'ids.npy'\n\ndef do_precompute(df):\n    if FEAT_A_FILE.exists() and FEAT_B_FILE.exists() and PROFILE_FILE.exists() and IDS_FILE.exists():\n        print(\"Precomputed files found. Loading.\")\n        feats_a = np.load(FEAT_A_FILE, mmap_mode='r')\n        feats_b = np.load(FEAT_B_FILE, mmap_mode='r')\n        profiles = np.load(PROFILE_FILE, mmap_mode='r')\n        ids = np.load(IDS_FILE)\n        return ids, feats_a, feats_b, profiles\n\n    print(\"Starting precompute of features and profiles ...\")\n    # setup backbone (frozen)\n    backbone = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)  # modern API\n    modules = list(backbone.children())[:-1]  # remove fc\n    backbone = nn.Sequential(*modules).to(DEVICE)\n    backbone.eval()\n    for p in backbone.parameters(): p.requires_grad = False\n\n    # image augment pipeline for precompute: two different strong-ish augmentations\n    aug1 = transforms.Compose([\n        transforms.ToPILImage(),\n        transforms.RandomResizedCrop(PRECOMP_IMG_SIZE, scale=(0.7,1.0), ratio=(0.9,1.1)),\n        transforms.RandomHorizontalFlip(),\n        transforms.ColorJitter(brightness=0.12, contrast=0.12, saturation=0.08),\n        transforms.ToTensor(),\n        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])\n    ])\n    aug2 = transforms.Compose([\n        transforms.ToPILImage(),\n        transforms.RandomResizedCrop(PRECOMP_IMG_SIZE, scale=(0.6,0.98), ratio=(0.85,1.15)),\n        transforms.RandomHorizontalFlip(),\n        transforms.RandomRotation(15),\n        transforms.ColorJitter(brightness=0.10, contrast=0.10, saturation=0.06),\n        transforms.ToTensor(),\n        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])\n    ])\n\n    def img_loader(path):\n        img = Image.open(path).convert('RGB')\n        return np.array(img)\n\n    N = len(df)\n    feats_a = np.zeros((N, BACKBONE_FEATURE_DIM), dtype=np.float32)\n    feats_b = np.zeros((N, BACKBONE_FEATURE_DIM), dtype=np.float32)\n    profiles = np.zeros((N, PROFILE_ANGLES), dtype=np.float32)\n    ids = []\n\n    # create indices list for batching\n    idxs = list(range(N))\n    batch_size = BATCH_SIZE_PRECOMP if torch.cuda.is_available() else 64\n\n    # helper: compute radial profile from resized image\n    def compute_profile_from_arr(arr_rgb, angles=PROFILE_ANGLES):\n        gray = cv2.cvtColor(arr_rgb, cv2.COLOR_RGB2GRAY)\n        blur = cv2.GaussianBlur(gray, (7,7), 0)\n        th = cv2.adaptiveThreshold(blur,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY_INV,11,2)\n        contours, _ = cv2.findContours(th, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n        if not contours:\n            return np.zeros(angles, dtype=np.float32)\n        c = max(contours, key=cv2.contourArea)\n        mask = np.zeros_like(th)\n        cv2.drawContours(mask, [c], -1, 255, -1)\n        kernel = np.ones((5,5), np.uint8)\n        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)\n        ys, xs = np.where(mask > 0)\n        if len(xs) == 0:\n            return np.zeros(angles, dtype=np.float32)\n        cx, cy = xs.mean(), ys.mean()\n        h, w = mask.shape\n        prof = np.zeros(angles, dtype=np.float32)\n        angles_arr = np.linspace(0, 2*np.pi, angles, endpoint=False)\n        for i, a in enumerate(angles_arr):\n            dx, dy = math.cos(a), math.sin(a)\n            r = 0\n            while True:\n                x = int(round(cx + r*dx)); y = int(round(cy + r*dy))\n                if x < 0 or x >= w or y < 0 or y >= h: break\n                if mask[y,x] == 0: break\n                r += 1\n                if r > max(w,h): break\n            prof[i] = r\n        mx = prof.max()\n        if mx > 0:\n            prof = prof / (mx + 1e-9)\n        return prof.astype(np.float32)\n\n    # iterate in batches to extract features (two augmented views)\n    with torch.no_grad():\n        for start in tqdm(range(0, N, batch_size), desc='Precompute batches'):\n            end = min(N, start + batch_size)\n            batch_indices = idxs[start:end]\n            # load images to list for augmentation\n            imgs = []\n            for i in batch_indices:\n                p = df.loc[i, 'image_path']\n                arr = img_loader(p)\n                imgs.append(arr)\n                ids.append(df.loc[i, 'id'])\n            # create two augmented batches\n            tensor_a = torch.stack([aug1(img) for img in imgs], dim=0).to(DEVICE)\n            tensor_b = torch.stack([aug2(img) for img in imgs], dim=0).to(DEVICE)\n            # backbone forward (B x C x H x W) -> global pool (B x feat)\n            featsA = backbone(tensor_a).view(len(batch_indices), -1).cpu().numpy()\n            featsB = backbone(tensor_b).view(len(batch_indices), -1).cpu().numpy()\n            # store\n            for j, idx in enumerate(batch_indices):\n                feats_a[idx, :] = featsA[j]\n                feats_b[idx, :] = featsB[j]\n            # compute profiles from resized (use PRECOMP_IMG_SIZE)\n            for j, arr in enumerate(imgs):\n                try:\n                    pil = Image.fromarray(arr).resize((PRECOMP_IMG_SIZE, PRECOMP_IMG_SIZE), Image.BICUBIC)\n                    arr_small = np.array(pil)\n                except Exception:\n                    arr_small = cv2.resize(arr, (PRECOMP_IMG_SIZE, PRECOMP_IMG_SIZE))\n                profiles[start + j] = compute_profile_from_arr(arr_small)\n    # save arrays to disk\n    np.save(FEAT_A_FILE, feats_a)\n    np.save(FEAT_B_FILE, feats_b)\n    np.save(PROFILE_FILE, profiles)\n    np.save(IDS_FILE, np.array(ids))\n    print(\"Precompute done. Files saved to:\", CACHE_ROOT)\n    return np.array(ids), feats_a, feats_b, profiles\n\nids, feats_a, feats_b, profiles = do_precompute(df)\n\n# -------------------------\n# 5) Create Dataset that loads precomputed feats and profiles\n# -------------------------\nclass PrecompContrastiveDataset(Dataset):\n    def __init__(self, feats_a, feats_b, profiles, ids):\n        assert len(feats_a) == len(feats_b) == len(profiles) == len(ids)\n        self.feats_a = feats_a.astype(np.float32)\n        self.feats_b = feats_b.astype(np.float32)\n        self.profiles = profiles.astype(np.float32)\n        self.ids = list(ids)\n\n    def __len__(self):\n        return len(self.ids)\n\n    def __getitem__(self, idx):\n        a = self.feats_a[idx]\n        b = self.feats_b[idx]\n        prof = self.profiles[idx]\n        return {\n            'id': self.ids[idx],\n            'a': torch.from_numpy(a),\n            'b': torch.from_numpy(b),\n            'profile': torch.from_numpy(prof)\n        }\n\n# -------------------------\n# 6) Head model that fuses feat + profile -> projection\n# -------------------------\nclass HeadNet(nn.Module):\n    def __init__(self, feat_dim=BACKBONE_FEATURE_DIM, profile_dim=PROFILE_ANGLES, proj_dim=PROJ_DIM):\n        super().__init__()\n        # small MLP for feat\n        self.feat_mlp = nn.Sequential(\n            nn.Linear(feat_dim, 1024),\n            nn.ReLU(),\n            nn.Linear(1024, proj_dim)\n        )\n        # small MLP for profile\n        self.prof_mlp = nn.Sequential(\n            nn.Linear(profile_dim, 256),\n            nn.ReLU(),\n            nn.Linear(256, proj_dim)\n        )\n        # final fusion projection\n        self.proj = nn.Sequential(\n            nn.Linear(proj_dim, proj_dim),\n            nn.ReLU(),\n            nn.Linear(proj_dim, proj_dim)\n        )\n    def forward(self, feat, prof):\n        # feat: B x feat_dim (not normalized); prof: B x profile_dim\n        f = self.feat_mlp(feat)\n        p = self.prof_mlp(prof)\n        f = F.normalize(f, dim=1)\n        p = F.normalize(p, dim=1)\n        fused = F.normalize(f + p, dim=1)  # B x proj_dim\n        z = self.proj(fused)\n        z = F.normalize(z, dim=1)\n        return z\n\n# -------------------------\n# 7) NT-Xent (efficient on precomputed embeddings)\n# -------------------------\ndef nt_xent(z1, z2, temperature=TEMPERATURE):\n    z = torch.cat([z1, z2], dim=0)  # 2N x D\n    sim = torch.matmul(z, z.t()) / temperature  # cosine since z normalized\n    N = z1.size(0)\n    mask = (~torch.eye(2*N, device=z.device).bool()).float()\n    exp_sim = torch.exp(sim) * mask\n    # positive sims: i <-> i+N\n    pos = torch.exp(torch.sum(z1 * z2, dim=1) / temperature)\n    pos = torch.cat([pos, pos], dim=0)\n    denom = exp_sim.sum(dim=1) + 1e-9\n    loss = -torch.log(pos / denom)\n    return loss.mean()\n\n# -------------------------\n# 8) Train head quickly\n# -------------------------\n# split dataset into train/val (90/10)\nN = len(ids)\nperm = np.random.RandomState(SEED).permutation(N)\nn_train = int(0.90 * N)\ntrain_idx = perm[:n_train]\nval_idx = perm[n_train:]\n\ntrain_dataset = PrecompContrastiveDataset(feats_a[train_idx], feats_b[train_idx], profiles[train_idx], ids[train_idx])\nval_dataset = PrecompContrastiveDataset(feats_a[val_idx], feats_b[val_idx], profiles[val_idx], ids[val_idx])\n\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE_TRAIN, shuffle=True, num_workers=NUM_WORKERS, pin_memory=True, drop_last=True)\nval_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE_TRAIN, shuffle=False, num_workers=0, pin_memory=True, drop_last=False)\n\nprint(\"Train/Val sizes:\", len(train_dataset), len(val_dataset))\n\nhead = HeadNet(feat_dim=BACKBONE_FEATURE_DIM, profile_dim=PROFILE_ANGLES, proj_dim=PROJ_DIM).to(DEVICE)\noptimizer = torch.optim.AdamW(head.parameters(), lr=LR_HEAD, weight_decay=1e-5)\nscaler = torch.cuda.amp.GradScaler() if torch.cuda.is_available() else None\n\nstart_time = time.time()\nfor epoch in range(1, EPOCHS_HEAD+1):\n    head.train()\n    total_loss = 0.0\n    count = 0\n    pbar = tqdm(train_loader, desc=f\"Epoch {epoch}/{EPOCHS_HEAD}\")\n    for batch in pbar:\n        a = batch['a'].to(DEVICE).float()\n        b = batch['b'].to(DEVICE).float()\n        prof = batch['profile'].to(DEVICE).float()\n        optimizer.zero_grad()\n        if scaler:\n            with torch.cuda.amp.autocast():\n                z1 = head(a, prof)\n                z2 = head(b, prof)\n                loss = nt_xent(z1, z2)\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n        else:\n            z1 = head(a, prof)\n            z2 = head(b, prof)\n            loss = nt_xent(z1, z2)\n            loss.backward()\n            optimizer.step()\n        batch_sz = a.size(0)\n        total_loss += float(loss.item()) * batch_sz\n        count += batch_sz\n        pbar.set_postfix(loss=total_loss / count)\n    avg_loss = total_loss / max(1, count)\n    print(f\"Epoch {epoch} train loss: {avg_loss:.5f}\")\n    # quick val step (compute average loss)\n    head.eval()\n    val_loss = 0.0; vcount = 0\n    with torch.no_grad():\n        for batch in val_loader:\n            a = batch['a'].to(DEVICE).float()\n            b = batch['b'].to(DEVICE).float()\n            prof = batch['profile'].to(DEVICE).float()\n            z1 = head(a, prof); z2 = head(b, prof)\n            loss = nt_xent(z1, z2)\n            val_loss += float(loss.item()) * a.size(0)\n            vcount += a.size(0)\n    if vcount:\n        print(f\"Val loss: {val_loss/vcount:.5f}\")\n    # save head checkpoint\n    ckpt = CACHE_ROOT / f'head_epoch{epoch}.pth'\n    torch.save(head.state_dict(), ckpt)\n    print(\"Saved checkpoint:\", ckpt)\n\nprint(\"Head training done in {:.1f} minutes\".format((time.time() - start_time)/60))\n\n# -------------------------\n# 9) Produce final embeddings for all images (fusion via head on backbone feats)\n# -------------------------\nprint(\"Producing final fused embeddings for all images ...\")\nall_dataset = PrecompContrastiveDataset(feats_a, feats_b, profiles, ids)  # feats_a used as canonical view\nall_loader = DataLoader(all_dataset, batch_size=256, shuffle=False, num_workers=0)\nhead.eval()\nembs = []\nwith torch.no_grad():\n    for batch in tqdm(all_loader):\n        f = batch['a'].to(DEVICE).float()\n        prof = batch['profile'].to(DEVICE).float()\n        z = head(f, prof)  # normalized\n        embs.append(z.cpu().numpy())\nemb_all = np.vstack(embs)\n# ensure normalized\nnorms = np.linalg.norm(emb_all, axis=1, keepdims=True)\nemb_all = emb_all / (norms + 1e-9)\nprint(\"Embeddings shape:\", emb_all.shape)\n\n# save embeddings\nnp.save(CACHE_ROOT / 'emb_all.npy', emb_all)\nnp.save(CACHE_ROOT / 'ids_all.npy', np.array(ids))\nprint(\"Saved final embeddings to cache\")\n\n# -------------------------\n# 10) Build top-k neighbors & assemble (fast)\n# -------------------------\nprint(\"Building top-k neighbors ...\")\nnn = NearestNeighbors(n_neighbors=min(TOP_K+1, len(emb_all)), metric='cosine', n_jobs=-1)\nnn.fit(emb_all)\ndists, idxs = nn.kneighbors(emb_all)\nneighbors = {}\nfor i, id_ in enumerate(ids):\n    neighs = []\n    for j_idx, d in zip(idxs[i], dists[i]):\n        if j_idx == i: continue\n        sim = 1.0 - float(d)\n        neighs.append((ids[j_idx], sim))\n        if len(neighs) >= TOP_K: break\n    neighbors[id_] = neighs\n\n# union-find assembly\nclass UF:\n    def __init__(self):\n        self.p = {}\n    def find(self, x):\n        if x not in self.p: self.p[x] = x; return x\n        while self.p[x] != x:\n            self.p[x] = self.p[self.p[x]]\n            x = self.p[x]\n        return x\n    def union(self, a, b):\n        ra, rb = self.find(a), self.find(b)\n        if ra != rb: self.p[rb] = ra\n    def comps(self):\n        groups = defaultdict(list)\n        for k in self.p.keys():\n            groups[self.find(k)].append(k)\n        return list(groups.values())\n\nuf = UF()\nfor n, neighs in neighbors.items():\n    uf.find(n)\n    for nid, sim in neighs:\n        if sim >= SIMILARITY_THRESHOLD:\n            uf.union(n, nid)\ncomponents = uf.comps()\n# fallback: if no unions (rare), make singletons\nif not components:\n    components = [[k] for k in neighbors.keys()]\n\nprint(\"Assembled components:\", len(components))\n# save components\nrows = []\nfor cid, comp in enumerate(components):\n    for s in comp:\n        rows.append({'component_id': cid, 'shard_id': s})\npd.DataFrame(rows).to_csv('/kaggle/working/assembled_components.csv', index=False)\nprint(\"Saved assembled components to /kaggle/working/assembled_components.csv\")\n\nprint(\"Pipeline complete. You can now iterate (augmentations, threshold, top_k, proj_dim) to improve score.\")\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-08-12T04:27:45.584278Z","iopub.execute_input":"2025-08-12T04:27:45.585290Z","iopub.status.idle":"2025-08-12T05:21:32.516902Z","shell.execute_reply.started":"2025-08-12T04:27:45.585252Z","shell.execute_reply":"2025-08-12T05:21:32.515963Z"}},"outputs":[],"execution_count":null}]}