{"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":"gpu","dataSources":[{"sourceType":"competition","sourceId":19989,"databundleVersionId":1160143}],"dockerImageVersionId":31090,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -q ultralytics","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-07T10:14:24.290647Z","iopub.execute_input":"2026-04-07T10:14:24.290939Z","iopub.status.idle":"2026-04-07T10:15:37.906511Z","shell.execute_reply.started":"2026-04-07T10:14:24.290900Z","shell.execute_reply":"2026-04-07T10:15:37.905626Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================\n# End-to-end pipeline: COCO→COCO, YOLO split, visualization,\n# DINOv2-guided self-supervised pretraining (distillation),\n# YOLO fine-tuning\n# =========================================================\n\nimport os, json, math, gc, random, shutil, ast\nfrom pathlib import Path\n\nimport torch, torch.nn as nn, torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nfrom PIL import Image\nimport numpy as np\nimport pandas as pd\nfrom sklearn.model_selection import train_test_split\nfrom tqdm.auto import tqdm\nimport matplotlib.pyplot as plt\n\nfrom ultralytics import YOLO\n\n# -----------------------\n# Paths & basic settings\n# -----------------------\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Device:\", device, \"| torch\", torch.__version__)\n\nBASE   = Path(\"/kaggle/input/global-wheat-detection\")\nWORK   = Path(\"/kaggle/working/wheat_ssl_yolov12_dinov2\")\nCOCO_S = WORK / \"coco_split\"\nYOLO_S = WORK / \"yolo_split\"\nDATA_Y = WORK / \"data_wheat.yaml\"  # YOLO data.yaml\nSSL_W  = WORK / \"backbone_ssl_v12_dinov2.pt\"   # conv+neck only (no Detect head)\nWORK.mkdir(parents=True, exist_ok=True)\n\n# For reproducibility (optional)\nrandom.seed(42)\nnp.random.seed(42)\ntorch.manual_seed(42)\nif device == \"cuda\":\n    torch.cuda.manual_seed_all(42)\n\n# =========================================================\n# 1) COCO → COCO split (train/val/test)\n# =========================================================\nprint(\"\\n[1/5] Building COCO→COCO split ...\")\nCOCO_S.mkdir(parents=True, exist_ok=True)\n\n# Load train.csv\ndf = pd.read_csv(BASE/\"train.csv\")\ndf[\"bbox\"] = df[\"bbox\"].apply(ast.literal_eval)  # \"[x,y,w,h]\" -> [x,y,w,h]\n\n# Split by unique image_id\ntrain_ids, val_ids = train_test_split(df[\"image_id\"].unique(), test_size=0.2, random_state=42)\ntrain_df = df[df[\"image_id\"].isin(train_ids)]\nval_df   = df[df[\"image_id\"].isin(val_ids)]\n\ndef build_coco(df_part, img_dir):\n    coco = {\n        \"images\": [],\n        \"annotations\": [],\n        \"categories\": [{\"id\": 1, \"name\": \"wheat_head\"}]\n    }\n    ann_id = 1\n    img_id_counter = 1\n    id_map = {}  # str image_id -> int id\n\n    for img_str, sub in tqdm(df_part.groupby(\"image_id\"), desc=\"COCO JSON\", unit=\"img\"):\n        img_path = img_dir / f\"{img_str}.jpg\"\n        if not img_path.exists():\n            continue\n\n        if img_str not in id_map:\n            id_map[img_str] = img_id_counter\n            img_id_counter += 1\n        img_id_num = id_map[img_str]\n\n        coco[\"images\"].append({\n            \"id\": img_id_num,\n            \"file_name\": f\"{img_str}.jpg\",\n            \"width\": int(sub.iloc[0][\"width\"]),\n            \"height\": int(sub.iloc[0][\"height\"])\n        })\n\n        for _, row in sub.iterrows():\n            x, y, w, h = row[\"bbox\"]\n            coco[\"annotations\"].append({\n                \"id\": ann_id,\n                \"image_id\": img_id_num,\n                \"category_id\": 1,\n                \"bbox\": [float(x), float(y), float(w), float(h)],\n                \"area\": float(w*h),\n                \"iscrowd\": 0\n            })\n            ann_id += 1\n\n    return coco\n\n# Save splits (COCO)\nfor split_name, split_df in [(\"train\", train_df), (\"val\", val_df)]:\n    out_dir = COCO_S / split_name\n    out_img = out_dir / \"images\"\n    out_dir.mkdir(parents=True, exist_ok=True)\n    out_img.mkdir(parents=True, exist_ok=True)\n\n    print(f\"  ➤ Copying {split_name} images (COCO) ...\")\n    for img_id in tqdm(split_df[\"image_id\"].unique(), unit=\"img\"):\n        src = BASE/\"train\"/f\"{img_id}.jpg\"\n        if src.exists():\n            shutil.copy(src, out_img/src.name)\n\n    print(f\"  ➤ Writing {split_name} annotations.json ...\")\n    coco_dict = build_coco(split_df, BASE/\"train\")\n    with open(out_dir/\"annotations.json\", \"w\") as f:\n        json.dump(coco_dict, f)\n\n# Test (unlabeled) for COCO split\nprint(\"  ➤ Copying test images (COCO) ...\")\n(COCO_S/\"test\"/\"images\").mkdir(parents=True, exist_ok=True)\nfor ext in (\"*.jpg\",\"*.jpeg\",\"*.png\"):\n    for img in tqdm((BASE/\"test\").glob(ext), unit=\"img\", desc=\"test\"):\n        shutil.copy(img, (COCO_S/\"test\"/\"images\")/img.name)\n\nprint(\"✓ COCO split ready at:\", COCO_S)\n\n# =========================================================\n# 2) YOLO-format split (train/val/test + data.yaml)\n#    Convert directly from original CSV (COCO→YOLO labels)\n# =========================================================\nprint(\"\\n[2/5] Building YOLO-format split ...\")\nYOLO_S.mkdir(parents=True, exist_ok=True)\n\ndef write_yolo_labels(df_part, out_im, out_lb, img_dir):\n    out_im.mkdir(parents=True, exist_ok=True)\n    out_lb.mkdir(parents=True, exist_ok=True)\n    # Create a set to avoid copying the same image multiple times\n    copied = set()\n    for _, row in df_part.iterrows():\n        x, y, w, h = row[\"bbox\"]\n        W, H = row[\"width\"], row[\"height\"]\n        # COCO -> YOLO normalized\n        xc = (x + w/2) / W\n        yc = (y + h/2) / H\n        wn = w / W\n        hn = h / H\n\n        label_path = out_lb / f\"{row['image_id']}.txt\"\n        with open(label_path, \"a\") as f:\n            f.write(f\"0 {xc:.6f} {yc:.6f} {wn:.6f} {hn:.6f}\\n\")\n\n        src = img_dir / f\"{row['image_id']}.jpg\"\n        if src.exists() and row['image_id'] not in copied:\n            shutil.copy(src, out_im/src.name)\n            copied.add(row['image_id'])\n\n# Train/Val\nprint(\"  ➤ Writing YOLO train labels+images ...\")\nwrite_yolo_labels(train_df, YOLO_S/\"train\"/\"images\", YOLO_S/\"train\"/\"labels\", BASE/\"train\")\n\nprint(\"  ➤ Writing YOLO val labels+images ...\")\nwrite_yolo_labels(val_df, YOLO_S/\"val\"/\"images\", YOLO_S/\"val\"/\"labels\", BASE/\"train\")\n\n# Test images\nprint(\"  ➤ Copying YOLO test images ...\")\n(YOLO_S/\"test\"/\"images\").mkdir(parents=True, exist_ok=True)\nfor ext in (\"*.jpg\",\"*.jpeg\",\"*.png\"):\n    for img in tqdm((BASE/\"test\").glob(ext), unit=\"img\", desc=\"test(yolo)\"):\n        shutil.copy(img, (YOLO_S/\"test\"/\"images\")/img.name)\n\n# YOLO data.yaml (simple \"yaml-like\" text)\ndata_yaml = {\n    \"path\": str(YOLO_S),\n    \"train\": \"train/images\",\n    \"val\": \"val/images\",\n    \"test\": \"test/images\",\n    \"nc\": 1,\n    \"names\": [\"wheat_head\"]\n}\nDATA_Y.write_text(\n    json.dumps(data_yaml, indent=2)\n    .replace(\"{\\n\", \"\")\n    .replace(\"}\", \"\")\n    .replace('\"', '')\n    .replace(\",\", \"\")\n)\n# If you prefer real yaml:\n# import yaml; DATA_Y.write_text(yaml.dump(data_yaml))\n\nprint(\"✓ YOLO split ready at:\", YOLO_S)\nprint(\"✓ data.yaml:\", DATA_Y)\n\n# =========================================================\n# 3) Visualize a random image with bboxes (COCO split)\n# =========================================================\nprint(\"\\n[3/5] Visualizing a random train image (COCO split) ...\")\nwith open(COCO_S/\"train\"/\"annotations.json\", \"r\") as f:\n    coco_train = json.load(f)\n\nif len(coco_train[\"images\"]) > 0:\n    img_info = random.choice(coco_train[\"images\"])\n    img_id_num = img_info[\"id\"]\n    img_file = COCO_S/\"train\"/\"images\"/img_info[\"file_name\"]\n\n    # Gather annots\n    anns = [a for a in coco_train[\"annotations\"] if a[\"image_id\"] == img_id_num]\n\n    # Draw\n    im = np.array(Image.open(img_file).convert(\"RGB\"))\n    for ann in anns:\n        x, y, w, h = ann[\"bbox\"]\n        x1, y1, x2, y2 = int(x), int(y), int(x+w), int(y+h)\n        im[y1:y1+2, x1:x2] = [255, 0, 0]\n        im[y2-2:y2, x1:x2] = [255, 0, 0]\n        im[y1:y2, x1:x1+2] = [255, 0, 0]\n        im[y1:y2, x2-2:x2] = [255, 0, 0]\n\n    plt.figure(figsize=(8,8))\n    plt.imshow(im)\n    plt.axis(\"off\")\n    plt.title(f\"{img_info['file_name']} — {len(anns)} boxes\")\n    plt.show()\nelse:\n    print(\"No images found in COCO train split to visualize.\")\n\n# =========================================================\n# 4) DINOv2-guided self-supervised pretraining on images\n#    (distill YOLO backbone features to DINOv2 CLS embeddings)\n# =========================================================\nprint(\"\\n[4/5] DINOv2-guided self-supervised pretraining ...\")\n\nMODEL = \"yolo12n.yaml\"     # YOLOv12 nano\nSSL_EPOCHS, SSL_BATCH = 5, 8  # adjust as you like\n\n# ---- Utilities / hooks ----\ndef get_detect_module(ultra_model):\n    # Ultralytics: detection module is last (Detect)\n    return ultra_model.model[-1]\n\nclass DetectInputHook:\n    \"\"\"Capture the pyramid features passed into Detect via forward_pre_hook.\"\"\"\n    def __init__(self, detect_module):\n        self.h = detect_module.register_forward_pre_hook(self._hook)\n        self.feats = None  # list of tensors [P3,P4,P5]\n    def _hook(self, module, args):\n        x = args[0]\n        if isinstance(x, (list, tuple)):\n            self.feats = [t.detach() for t in x]\n        else:\n            self.feats = [x.detach()]\n    def close(self):\n        self.h.remove()\n\ndef global_pool_concat(feats):\n    # list [N,C,H,W] -> [N, sumC]\n    pooled = [F.adaptive_avg_pool2d(f, 1).flatten(1) for f in feats]\n    return torch.cat(pooled, dim=1)\n\n# Two-view SSL dataset\nclass TwoView(Dataset):\n    def __init__(self, img_dirs, augment):\n        self.paths = []\n        for d in img_dirs:\n            self.paths += list(Path(d).glob(\"*.jpg\")) + list(Path(d).glob(\"*.jpeg\")) + list(Path(d).glob(\"*.png\"))\n        self.augment = augment\n    def __len__(self): return len(self.paths)\n    def __getitem__(self, idx):\n        p = self.paths[idx]\n        img = Image.open(p).convert(\"RGB\")\n        v1 = self.augment(img); v2 = self.augment(img)\n        return v1, v2\n\n# SSL augment pipeline\naugment = transforms.Compose([\n    transforms.RandomResizedCrop(224, scale=(0.2, 1.0)),  # 224 aligns with ViT\n    transforms.RandomHorizontalFlip(p=0.5),\n    transforms.RandomApply([transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),\n    transforms.RandomGrayscale(p=0.2),\n    transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),\n    transforms.ToTensor()\n])\n\n# Image dirs for SSL\ntrain_img_dir = YOLO_S / \"train\" / \"images\"\nval_img_dir   = YOLO_S / \"val\"   / \"images\"\nds = TwoView([train_img_dir, val_img_dir], augment)\ndl = DataLoader(ds, batch_size=SSL_BATCH, shuffle=True, num_workers=0, pin_memory=False, drop_last=True)\n\n# ---------- DINOv2 teacher (frozen) ----------\ndef load_dinov2_teacher(device=\"cpu\"):\n    teacher = None\n    try:\n        # Preferred (official) loader\n        teacher = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14')\n    except Exception:\n        try:\n            import timm\n            teacher = timm.create_model('vit_base_patch14_dinov2', pretrained=True)\n        except Exception as e:\n            raise RuntimeError(\n                \"Failed to load DINOv2. Enable internet for torch.hub the first time, \"\n                \"or install timm with a dinov2 model.\"\n            ) from e\n    teacher.eval().to(device)\n\n    # Wrap to return CLS token (dim=768 for ViT-B/14)\n    class DINOv2CLS(torch.nn.Module):\n        def __init__(self, backbone):\n            super().__init__()\n            self.backbone = backbone\n        @torch.no_grad()\n        def forward(self, x):\n            out = self.backbone(x)\n            # Handle dict outputs (official dinov2 exposes x_norm_clstoken)\n            if isinstance(out, dict):\n                for k in [\"x_norm_clstoken\", \"cls_token\", \"last_hidden_state\", \"head\"]:\n                    if k in out:\n                        out = out[k]\n                        break\n            # If shape is [N, C], assume CLS already\n            if isinstance(out, torch.Tensor) and out.ndim == 2:\n                return out  # [N, 768]\n            # If shape is [N, T, C], take CLS at index 0\n            if isinstance(out, torch.Tensor) and out.ndim == 3 and out.shape[1] >= 1:\n                return out[:, 0, :]\n            # Fallbacks\n            if isinstance(out, torch.Tensor) and out.ndim == 3:\n                return out.mean(dim=1)\n            if isinstance(out, torch.Tensor) and out.ndim == 4:\n                return F.adaptive_avg_pool2d(out, 1).flatten(1)\n            return out\n    return DINOv2CLS(teacher)\n\n# ---------- Build/Load YOLO backbone and DINOv2 ----------\nif SSL_W.exists():\n    print(\"✓ DINOv2-distilled backbone cached – skipping pre-train\")\nelse:\n    print(\"🔄  Running DINOv2-guided pre-training …\")\n    # YOLO encoder (train conv+neck; Detect head excluded when saving)\n    online_full = YOLO(MODEL).model.to(device)\n    online_hook = DetectInputHook(get_detect_module(online_full))\n\n    # DINOv2 teacher (frozen)\n    dinov2 = load_dinov2_teacher(device=device)\n    for p in dinov2.parameters():\n        p.requires_grad = False\n\n    # Infer YOLO pooled feature dim\n    with torch.no_grad():\n        dmy = torch.zeros(1,3,224,224,device=device)\n        _ = online_full(dmy)\n        feat_dim = global_pool_concat(online_hook.feats).shape[1]\n    print(\"YOLO pooled feat dim:\", feat_dim, \"| DINOv2 CLS dim: 768\")\n\n    # Projector: YOLO_feats -> 768 (match DINOv2 CLS)\n    projector = nn.Sequential(\n        nn.Linear(feat_dim, 1024, bias=False),\n        nn.BatchNorm1d(1024), nn.ReLU(inplace=True),\n        nn.Linear(1024, 768, bias=True)\n    ).to(device)\n\n    # Optimizer\n    opt = torch.optim.AdamW(\n        list(online_full.parameters()) + list(projector.parameters()),\n        lr=1e-3, weight_decay=1e-4\n    )\n    scaler = torch.cuda.amp.GradScaler(enabled=(device==\"cuda\"))\n\n    # Losses\n    mse = nn.MSELoss()\n    def cosine_loss(a, b, eps=1e-8):\n        a = F.normalize(a, dim=1)\n        b = F.normalize(b, dim=1)\n        return 1 - (a*b).sum(dim=1).mean()\n\n    for ep in range(SSL_EPOCHS):\n        online_full.train(); projector.train()\n        epoch_l, pbar = 0.0, tqdm(dl, desc=f\"DINOv2 distill {ep+1}/{SSL_EPOCHS}\", leave=False)\n        for v1, v2 in pbar:\n            v1, v2 = v1.to(device, non_blocking=True), v2.to(device, non_blocking=True)\n\n            with torch.autocast(device_type='cuda', enabled=(device==\"cuda\")):\n                # YOLO features (hook captures P3-P5)\n                _ = online_full(v1); y1 = global_pool_concat(online_hook.feats)\n                _ = online_full(v2); y2 = global_pool_concat(online_hook.feats)\n\n                # DINOv2 CLS (teacher)\n                with torch.no_grad():\n                    t1 = dinov2(v1)  # [N,768]\n                    t2 = dinov2(v2)  # [N,768]\n\n                # Project YOLO to 768\n                p1 = projector(y1)\n                p2 = projector(y2)\n\n                # Distillation: cross-view matching + mild view consistency\n                loss_distill = mse(p1, t2) + mse(p2, t1)\n                loss_cons    = 0.2 * cosine_loss(p1, p2)\n                loss = loss_distill + loss_cons\n\n            scaler.scale(loss).backward()\n            scaler.step(opt); scaler.update(); opt.zero_grad(set_to_none=True)\n            epoch_l += loss.item()\n            pbar.set_postfix(loss=f\"{loss.item():.4f}\")\n\n        print(f\"Epoch {ep+1}/{SSL_EPOCHS} DINOv2 distill loss={epoch_l/len(dl):.4f}\")\n\n    # Save ONLY backbone+neck weights (exclude Detect head)\n    sd = online_full.state_dict()\n    detect_mod = len(online_full.model) - 1\n    filtered = {k:v for k,v in sd.items() if not k.startswith(f\"model.{detect_mod}.\")}\n    torch.save(filtered, SSL_W)\n\n    online_hook.close()\n    del online_full, projector, dinov2, dl, ds\n    gc.collect()\n    if device == \"cuda\": torch.cuda.empty_cache()\n    print(\"✓ Saved SSL backbone weights →\", SSL_W)\n\n# =========================================================\n# 5) Supervised fine-tuning with YOLO using SSL backbone\n# =========================================================\n\n","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}]}