Initial commit with Git LFS for binaries

.gitattributes

@@ -0,0 +1,3 @@
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.JPEG filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,8 @@
+__pycache__/
+*.py[cod]
+*.class
+.DS_Store
+checkpoints/
+=4.0.0
+backbone_keys.txt
+keys.txt

README.md

@@ -0,0 +1,28 @@
+---
+title: GeoViG — Geometry-Aware Graph Vision
+emoji: 🔷
+colorFrom: indigo
+colorTo: blue
+sdk: gradio
+sdk_version: "4.44.0"
+app_file: app.py
+pre_load_timeout: 300
+pinned: true
+license: apache-2.0
+short_description: ImageNet + detection + medical segmentation with GeoViG
+---
+
+# GeoViG Gradio Demo
+
+Interactive demo for **GeoViG** — a lightweight hybrid CNN-GNN architecture featuring **SpreadEdgePool** and **GraphMRConv**.
+
+### 🌟 Features
+- **🖼️ Image Classification**: ImageNet-1K predictions with Ti/S/M/B variants.
+- **📦 Object Detection**: Mask R-CNN outputs for COCO objects.
+- **🏥 Medical Segmentation**: Polyp (Kvasir-SEG) and Nuclei (DSB 2018) segmentation.
+
+### 🛠️ Local Setup
+To run this locally:
+1. `pip install -r requirements.txt`
+2. `sh pre_run.sh` (to install MMCV/MMDet)
+3. `python app.py`

app.py

@@ -0,0 +1,442 @@
+"""
+GeoViG Multi-Task Gradio Demo
+Tabs:
+  1. Image Classification (ImageNet-1K)  — Ti / S / M / B
+  2. Object Detection & Segmentation     — GeoViG-M on COCO 2017
+  3. Medical Segmentation                — GeoViG-M on Kvasir-SEG / DSB 2018
+"""
+
+import sys, os, json, urllib.request
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+from PIL import Image, ImageDraw
+from huggingface_hub import hf_hub_download
+import numpy as np
+import gradio as gr
+
+# Detect Gradio API version — theme/css moved to launch() in v6
+_GR_MAJOR = int(gr.__version__.split(".")[0])
+_BLOCKS_KWARGS = {} if _GR_MAJOR >= 6 else {
+    "theme": gr.themes.Soft(primary_hue="indigo"),
+    "css": "footer { display: none !important; }"
+}
+_LAUNCH_KWARGS = {
+    "theme": gr.themes.Soft(primary_hue="indigo"),
+    "css": "footer { display: none !important; }"
+} if _GR_MAJOR >= 6 else {}
+
+from geovig import geovig_ti, geovig_s, geovig_m, geovig_b
+
+# ── Paths ──────────────────────────────────────────────────────────────────────
+BASE = "."
+REPO_ID = "OmarAlasqa/GeoViG"
+CKPT_FILES = {
+    "GeoViG-Ti": "pth/geovig_ti_5e4_8G_300_75_22/checkpoint.pth",
+    "GeoViG-S":  "pth/geovig_s_5e4_8G_300_77_48/checkpoint.pth",
+    "GeoViG-M":  "pth/geovig_m_5e4_8G_300_80_70/checkpoint.pth",
+    "GeoViG-B":  "pth/geovig_b_5e4_8G_300_82_38/checkpoint.pth",
+    "det_m":     "coco_det_seg_pth/geovig_m_det_seg/epoch_12.pth",
+    "det_b":     "coco_det_seg_pth/geovig_b_det_seg/epoch_12.pth",
+    "kvasir_m":  "medical/kvasir_geovig_m/checkpoint.pth",
+    "dsb_m":     "medical/dsb_geovig_m/checkpoint.pth"
+}
+
+def get_ckpt_path(key):
+    return hf_hub_download(repo_id=REPO_ID, filename=CKPT_FILES[key])
+
+MMDET_DIR = "."
+
+# ── ImageNet labels ────────────────────────────────────────────────────────────
+LABELS_FILE = "imagenet_labels.json"
+if not os.path.exists(LABELS_FILE):
+    urllib.request.urlretrieve(
+        "https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json",
+        LABELS_FILE)
+with open(LABELS_FILE) as f:
+    IMAGENET_LABELS = json.load(f)
+
+# ── MMDetection availability ───────────────────────────────────────────────────
+try:
+    sys.path.insert(0, MMDET_DIR)
+    import mmcv
+    from mmdet.apis import init_detector, inference_detector
+    import mmdet.models  # noqa – registers all standard detectors
+    import geovig_det_backbone  # noqa – local copy: registers geovig_m_feat / geovig_b_feat
+    MMDET_OK = True
+    MMDET_ERROR = ""
+except Exception as e:
+    MMDET_OK = False
+    MMDET_ERROR = str(e)
+
+# ── Pre-processing ─────────────────────────────────────────────────────────────
+IMAGENET_TF = transforms.Compose([
+    transforms.Resize(256, interpolation=transforms.InterpolationMode.BICUBIC),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+MEDICAL_TF = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+# ── Model cache ────────────────────────────────────────────────────────────────
+MODEL_CACHE = {}
+
+BUILDERS = {
+    "GeoViG-Ti": geovig_ti,
+    "GeoViG-S":  geovig_s,
+    "GeoViG-M":  geovig_m,
+    "GeoViG-B":  geovig_b,
+}
+
+def _load_cls_model(variant: str):
+    if variant in MODEL_CACHE:
+        return MODEL_CACHE[variant]
+    model = BUILDERS[variant]()
+    ckpt_path = get_ckpt_path(variant)
+    try:
+        sd = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    except TypeError:
+        sd = torch.load(ckpt_path, map_location="cpu")  # PyTorch < 1.13
+    sd = sd.get("model", sd.get("state_dict", sd))
+    model.load_state_dict(sd, strict=False)
+    model.eval()
+    MODEL_CACHE[variant] = model
+    return model
+
+# ── TAB 1 — Classification ─────────────────────────────────────────────────────
+def classify(image: Image.Image, variant: str, top_k: int):
+    if image is None:
+        return {}
+    model = _load_cls_model(variant)
+    x = IMAGENET_TF(image.convert("RGB")).unsqueeze(0)
+    with torch.no_grad():
+        probs = F.softmax(model(x), dim=-1)[0]
+    top_probs, top_idxs = probs.topk(int(top_k))
+    return {IMAGENET_LABELS[i.item()]: float(p) for p, i in zip(top_probs, top_idxs)}
+
+# ── TAB 2 — Detection (COCO) ──────────────────────────────────────────────────
+COCO_NAMES = [
+    "person","bicycle","car","motorcycle","airplane","bus","train","truck","boat",
+    "traffic light","fire hydrant","stop sign","parking meter","bench","bird","cat",
+    "dog","horse","sheep","cow","elephant","bear","zebra","giraffe","backpack",
+    "umbrella","handbag","tie","suitcase","frisbee","skis","snowboard","sports ball",
+    "kite","baseball bat","baseball glove","skateboard","surfboard","tennis racket",
+    "bottle","wine glass","cup","fork","knife","spoon","bowl","banana","apple",
+    "sandwich","orange","broccoli","carrot","hot dog","pizza","donut","cake","chair",
+    "couch","potted plant","bed","dining table","toilet","tv","laptop","mouse",
+    "remote","keyboard","cell phone","microwave","oven","toaster","sink","refrigerator",
+    "book","clock","vase","scissors","teddy bear","hair drier","toothbrush"
+]
+
+def _draw_detections(image: Image.Image, result, score_thr=0.3):
+    """Draw bboxes and masks from MMDet result onto PIL image."""
+    img_draw = image.convert("RGBA")
+    draw = ImageDraw.Draw(img_draw)
+
+    # For the separate segmentation image
+    seg_img_draw = image.convert("RGBA")
+    seg_draw = ImageDraw.Draw(seg_img_draw)
+    seg_np_img = np.array(image.convert("RGB"))
+
+    if isinstance(result, tuple):
+        bbox_result, segm_result = result
+    else:
+        bbox_result, segm_result = result, None
+
+    for cls_id, bboxes in enumerate(bbox_result):
+        color = tuple(np.random.randint(100, 255, 3).tolist())
+        for i, bbox in enumerate(bboxes):
+            if bbox[4] < score_thr:
+                continue
+            x1, y1, x2, y2, score = bbox
+            draw.rectangle([x1, y1, x2, y2], outline=color, width=4)
+            label = f"{COCO_NAMES[cls_id]} {score:.2f}"
+            # Draw text with a small black outline for readability
+            for offset in [(-1,-1), (1,-1), (-1,1), (1,1)]:
+                draw.text((x1 + 2 + offset[0], y1 + 2 + offset[1]), label, fill=(0, 0, 0))
+            draw.text((x1 + 2, y1 + 2), label, fill=(255, 255, 0))
+
+            # Draw mask if available
+            if segm_result is not None and len(segm_result) > cls_id and segm_result[cls_id] is not None:
+                try:
+                    mask = segm_result[cls_id][i]
+                    if isinstance(mask, dict):
+                        import pycocotools.mask as mask_util
+                        mask = mask_util.decode(mask)
+                    
+                    # Ensure mask is boolean and same shape
+                    mask_bool = mask.astype(bool)
+                    # Alpha blending ONLY on the masked pixels
+                    # Background (1 - alpha) * image + alpha * color
+                    alpha = 0.5
+                    img_pixels = seg_np_img[mask_bool].astype(float)
+                    color_pixels = np.array(color, dtype=float)
+                    blended = (img_pixels * (1 - alpha) + color_pixels * alpha).astype(np.uint8)
+                    seg_np_img[mask_bool] = blended
+                except Exception:
+                    pass
+
+    det_result = img_draw.convert("RGB")
+    seg_result = Image.fromarray(seg_np_img)
+    return det_result, seg_result
+
+DET_MODEL_CACHE = {}
+
+def detect(image: Image.Image, model_size: str, score_thr: float):
+    if image is None:
+        return None, None, "⚠️ Please upload an image."
+    if not MMDET_OK:
+        return None, None, f"⚠️ MMDetection not available.\nInstall: `pip install mmcv-full==1.7.2 mmdet==2.28.0`\nError: {MMDET_ERROR}"
+
+    key = f"det_{model_size.lower()}"
+    if key not in DET_MODEL_CACHE:
+        cfg_map = {
+            "M": f"{MMDET_DIR}/configs/mask_rcnn_geovig_m_fpn_1x_coco.py",
+            "B": f"{MMDET_DIR}/configs/mask_rcnn_geovig_b_fpn_1x_coco.py",
+        }
+        ckpt_path = get_ckpt_path(f"det_{model_size.lower()}")
+        try:
+            det_model = init_detector(cfg_map[model_size], ckpt_path, device="cpu")
+            DET_MODEL_CACHE[key] = det_model
+        except Exception as e:
+            return None, None, f"⚠️ Failed to load detection model:\n{e}"
+
+    det_model = DET_MODEL_CACHE[key]
+    import cv2
+    cv_img = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
+    try:
+        result = inference_detector(det_model, cv_img)
+        bbox_result = result[0] if isinstance(result, tuple) else result
+        max_score = max([b[:, 4].max() if b.size > 0 else 0 for b in bbox_result])
+        print(f"DEBUG [{model_size} COCO]: Max Score = {max_score:.4f}")
+        
+        det_img, seg_img = _draw_detections(image, result, score_thr=score_thr)
+        
+        # Count detections
+        count = sum(len([b for b in cls_bboxes if b[4] >= score_thr]) for cls_bboxes in result[0])
+        status = f"✅ GeoViG-{model_size} detection complete. Found {count} objects."
+        if count == 0:
+            status += " (Try lowering score threshold)"
+            
+        return det_img, seg_img, status
+    except Exception as e:
+        return None, None, f"⚠️ Inference error:\n{e}"
+
+# ── TAB 3 — Medical Segmentation ──────────────────────────────────────────────
+MED_MODEL_CACHE = {}
+
+def _overlay_mask(image: Image.Image, mask: np.ndarray, color=(0, 200, 100), alpha=0.7):
+    """Overlay a binary mask on the PIL image."""
+    img_arr = np.array(image.convert("RGB")).astype(float)
+    overlay = img_arr.copy()
+    overlay[mask > 0.5] = np.array(color, dtype=float)
+    blended = (img_arr * (1 - alpha) + overlay * alpha).astype(np.uint8)
+    return Image.fromarray(blended)
+
+def segment_medical(image: Image.Image, dataset: str, score_thr: float):
+    if image is None:
+        return None, "⚠️ Please upload an image."
+    if not MMDET_OK:
+        return None, f"⚠️ MMDetection not available.\nInstall: `pip install mmcv-full==1.7.2 mmdet==2.28.0`\nError: {MMDET_ERROR}"
+
+    key = f"med_{dataset.lower().replace(' ', '_')}"
+    if key not in MED_MODEL_CACHE:
+        cfg_map = {
+            "Kvasir-SEG": f"{MMDET_DIR}/configs/kvasir/mask_rcnn_geovig_m_fpn_1x_kvasir.py",
+            "DSB 2018":   f"{MMDET_DIR}/configs/dsb/mask_rcnn_geovig_m_fpn_1x_dsb.py",
+        }
+        v_key = "kvasir_m" if dataset == "Kvasir-SEG" else "dsb_m"
+        ckpt_path = get_ckpt_path(v_key)
+        try:
+            med_model = init_detector(cfg_map[dataset], ckpt_path, device="cpu")
+            MED_MODEL_CACHE[key] = med_model
+        except Exception as e:
+            return None, f"⚠️ Failed to load model for {dataset}:\n{e}"
+
+    med_model = MED_MODEL_CACHE[key]
+    import cv2
+    cv_img = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
+    try:
+        result = inference_detector(med_model, cv_img)
+        bbox_result = result[0] if isinstance(result, tuple) else result
+        segm_result = result[1] if isinstance(result, tuple) else None
+        
+        max_score = max([b[:, 4].max() if b.size > 0 else 0 for b in bbox_result])
+        print(f"DEBUG [{dataset}]: Max Score = {max_score:.4f}")
+
+        # Merge all masks above threshold
+        combined_mask = np.zeros(cv_img.shape[:2], dtype=np.float32)
+        count = 0
+        if segm_result:
+            for cls_id, cls_masks in enumerate(segm_result):
+                bboxes = bbox_result[cls_id]
+                for i, mask in enumerate(cls_masks):
+                    if bboxes[i, 4] < score_thr:
+                        continue
+                    if isinstance(mask, dict):
+                        import pycocotools.mask as mask_util
+                        mask = mask_util.decode(mask)
+                    combined_mask = np.maximum(combined_mask, mask.astype(np.float32))
+                    count += 1
+
+        if count == 0:
+            return image, f"⚠️ No objects detected above score {score_thr}."
+        
+        color = (0, 200, 100) if dataset == "Kvasir-SEG" else (100, 150, 255)
+        out_img = _overlay_mask(image, combined_mask, color=color)
+        
+        # Now draw boxes on top of the masked image
+        draw = ImageDraw.Draw(out_img)
+        if bbox_result:
+            for cls_bboxes in bbox_result:
+                for bbox in cls_bboxes:
+                    if bbox[4] >= score_thr:
+                        x1, y1, x2, y2, score = bbox
+                        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
+
+        return out_img, f"✅ GeoViG-M medical segmentation complete. Found {count} instances."
+    except Exception as e:
+        return None, f"⚠️ Medical inference error:\n{e}"
+
+# ── Gradio Layout ──────────────────────────────────────────────────────────────
+with gr.Blocks(title="GeoViG Demo", **_BLOCKS_KWARGS) as demo:
+
+    gr.HTML("""
+    <div style="text-align:center; padding: 1rem 0 0.5rem">
+      <h1 style="font-size:2rem; font-weight:700; margin:0">🔷 GeoViG</h1>
+      <p style="color:#6b7280; margin:0.3rem 0 0">
+        Geometry-Aware Graph Reasoning for Mobile Vision &nbsp;|&nbsp;
+        <a href="https://huggingface.co/OmarAlasqa/GeoViG" target="_blank">Model Card</a>
+      </p>
+    </div>
+    """)
+
+    with gr.Tabs():
+
+        # ── Tab 1: Classification ──────────────────────────────────────────────
+        with gr.TabItem("🖼️ Image Classification"):
+            gr.Markdown("**ImageNet-1K Classification** — Upload any image and get top-K predictions.")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    cls_image   = gr.Image(type="pil", label="Input Image", height=280)
+                    cls_variant = gr.Radio(
+                        choices=["GeoViG-Ti", "GeoViG-S", "GeoViG-M", "GeoViG-B"],
+                        value="GeoViG-Ti",
+                        label="Model",
+                        info="Ti 75.2% · S 77.5% · M 80.7% · B 82.4% (ImageNet Top-1)"
+                    )
+                    cls_topk  = gr.Slider(1, 10, value=5, step=1, label="Top-K")
+                    cls_btn   = gr.Button("🚀 Classify", variant="primary")
+                with gr.Column(scale=1):
+                    cls_output = gr.Label(label="Predictions", num_top_classes=10)
+
+            cls_btn.click(fn=classify, inputs=[cls_image, cls_variant, cls_topk], outputs=cls_output)
+            cls_image.change(fn=classify, inputs=[cls_image, cls_variant, cls_topk], outputs=cls_output)
+
+            gr.Examples(
+                examples=[
+                    ["examples/n03394916_14162.JPEG", "GeoViG-Ti", 5],
+                    ["examples/n03417042_2960.JPEG", "GeoViG-M", 5],
+                ],
+                inputs=[cls_image, cls_variant, cls_topk],
+                outputs=cls_output,
+                fn=classify,
+                cache_examples=False,
+            )
+
+            gr.Markdown("""
+            | Model | Params | MACs | Top-1 | GPU (iPhone 13) |
+            |:---|:---:|:---:|:---:|:---:|
+            | GeoViG-Ti | 3.5M | 0.9G | 75.2% | 10.1ms |
+            | GeoViG-S  | 5.0M | 1.2G | 77.5% | 12.2ms |
+            | GeoViG-M  | 10.3M | 2.2G | 80.7% | 14.4ms |
+            | GeoViG-B  | 19.7M | 4.5G | 82.4% | 29.0ms |
+            """)
+
+        # ── Tab 2: Object Detection ────────────────────────────────────────────
+        with gr.TabItem("📦 Object Detection (COCO)"):
+            mmdet_status = "✅ MMDetection available" if MMDET_OK else "⚠️ MMDetection not installed — install `mmcv-full==1.7.2` and `mmdet==2.28.0` to enable this tab"
+            gr.Markdown(f"**MS COCO Object Detection & Instance Segmentation** (Mask R-CNN backbone)\n\n_{mmdet_status}_")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    det_image = gr.Image(type="pil", label="Input Image", height=280)
+                    det_size  = gr.Radio(choices=["M", "B"], value="M",
+                                         label="Backbone",
+                                         info="M: 40.7 Box AP · B: 42.5 Box AP")
+                    det_thr   = gr.Slider(0.1, 0.9, value=0.7, step=0.05, label="Score threshold")
+                    det_btn   = gr.Button("🔍 Detect & Segment", variant="primary", interactive=MMDET_OK)
+                    det_output_txt = gr.Textbox(label="Status", lines=2)
+                with gr.Column(scale=2):
+                    with gr.Row():
+                        det_output_img = gr.Image(label="Detection (Boxes)", type="pil")
+                        seg_output_img = gr.Image(label="Instance Segmentation (Masks)", type="pil")
+
+            det_btn.click(fn=detect, inputs=[det_image, det_size, det_thr],
+                          outputs=[det_output_img, seg_output_img, det_output_txt])
+
+            gr.Examples(
+                examples=[
+                    ["examples/000000032081.jpg", "M", 0.7],
+                    ["examples/000000045472.jpg", "M", 0.7],
+                    ["examples/cat.png", "B", 0.7],
+                ],
+                inputs=[det_image, det_size, det_thr],
+                outputs=[det_output_img, seg_output_img, det_output_txt],
+                fn=detect,
+                cache_examples=False,
+            )
+
+            gr.Markdown("""
+            | Backbone | Params | Box AP | Mask AP |
+            |:---|:---:|:---:|:---:|
+            | GeoViG-M | 10.3M | 40.7 | 37.7 |
+            | GeoViG-B | 19.7M | 42.5 | 38.9 |
+            """)
+
+        # ── Tab 3: Medical Segmentation ────────────────────────────────────────
+        with gr.TabItem("🏥 Medical Segmentation"):
+            gr.Markdown(f"**Medical Image Segmentation** — GeoViG-M encoder · Mask R-CNN decoder\n\n_{mmdet_status}_")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    med_image   = gr.Image(type="pil", label="Medical Image", height=280)
+                    med_dataset = gr.Radio(
+                        choices=["Kvasir-SEG", "DSB 2018"],
+                        value="Kvasir-SEG",
+                        label="Dataset / Task",
+                        info="Kvasir: Polyp segmentation (Dice 0.945, HD 12.94) · DSB: Nuclei (Dice 0.908)"
+                    )
+                    med_thr   = gr.Slider(0.1, 0.9, value=0.7, step=0.05, label="Score threshold")
+                    med_btn   = gr.Button("🔬 Segment", variant="primary", interactive=MMDET_OK)
+                with gr.Column(scale=1):
+                    med_output_img = gr.Image(label="Segmentation Result", type="pil")
+                    med_output_txt = gr.Textbox(label="Status", lines=3)
+
+            med_btn.click(fn=segment_medical, inputs=[med_image, med_dataset, med_thr],
+                          outputs=[med_output_img, med_output_txt])
+            
+            gr.Examples(
+                examples=[
+                    ["examples/cju5y84q3mdv50817eyp82xf3.jpg", "Kvasir-SEG", 0.7],
+                    ["examples/0c2550a23b8a0f29a7575de8c61690d3c31bc897dd5ba66caec201d201a278c2.png", "DSB 2018", 0.7],
+                ],
+                inputs=[med_image, med_dataset, med_thr],
+                outputs=[med_output_img, med_output_txt],
+                fn=segment_medical,
+                cache_examples=False,
+            )
+
+            gr.Markdown("""
+            | Task | Dataset | Dice ↑ | IoU ↑ | Hausdorff ↓ |
+            |:---|:---|:---:|:---:|:---:|
+            | Polyp Seg | Kvasir-SEG | **0.945** | **0.909** | **12.94** |
+            | Nuclei Seg | DSB 2018 | **0.908** | **0.839** | **5.19** |
+
+            > GeoViG reduces Hausdorff Distance **5.4× vs ResNet-50** (70.37 → 12.94) on Kvasir-SEG.
+            """)
+
+demo.launch(**_LAUNCH_KWARGS)

configs/_base_/datasets/coco_detection.py

@@ -0,0 +1,49 @@
+# dataset settings
+dataset_type = 'CocoDataset'
+data_root = '../../mmdetection/data/coco/'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations', with_bbox=True),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=2,
+    workers_per_gpu=2,
+    train=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_train2017.json',
+        img_prefix=data_root + 'train2017/',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_val2017.json',
+        img_prefix=data_root + 'val2017/',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_val2017.json',
+        img_prefix=data_root + 'val2017/',
+        pipeline=test_pipeline))
+evaluation = dict(interval=1, metric='bbox')

configs/_base_/datasets/coco_instance.py

@@ -0,0 +1,49 @@
+# dataset settings
+dataset_type = 'CocoDataset'
+data_root = '../../../hvig/detection/data/coco/'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations', with_bbox=True, with_mask=True),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=1, # was 2, made 1 to avoid out of memory
+    workers_per_gpu=2,
+    train=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_train2017.json',
+        img_prefix=data_root + 'train2017/',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_val2017.json',
+        img_prefix=data_root + 'val2017/',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_val2017.json',
+        img_prefix=data_root + 'val2017/',
+        pipeline=test_pipeline))
+evaluation = dict(metric=['bbox', 'segm'])

configs/_base_/datasets/coco_instance_semantic.py

@@ -0,0 +1,54 @@
+# dataset settings
+dataset_type = 'CocoDataset'
+data_root = 'data/coco/'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='LoadAnnotations', with_bbox=True, with_mask=True, with_seg=True),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='SegRescale', scale_factor=1 / 8),
+    dict(type='DefaultFormatBundle'),
+    dict(
+        type='Collect',
+        keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks', 'gt_semantic_seg']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip', flip_ratio=0.5),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=2,
+    workers_per_gpu=2,
+    train=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_train2017.json',
+        img_prefix=data_root + 'train2017/',
+        seg_prefix=data_root + 'stuffthingmaps/train2017/',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_val2017.json',
+        img_prefix=data_root + 'val2017/',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_val2017.json',
+        img_prefix=data_root + 'val2017/',
+        pipeline=test_pipeline))
+evaluation = dict(metric=['bbox', 'segm'])

configs/_base_/default_runtime.py

@@ -0,0 +1,16 @@
+checkpoint_config = dict(interval=1)
+# yapf:disable
+log_config = dict(
+    interval=50,
+    hooks=[
+        dict(type='TextLoggerHook'),
+        # dict(type='TensorboardLoggerHook')
+    ])
+# yapf:enable
+custom_hooks = [dict(type='NumClassCheckHook')]
+
+dist_params = dict(backend='nccl')
+log_level = 'INFO'
+load_from = None
+resume_from = None
+workflow = [('train', 1)]

configs/_base_/models/mask_rcnn_r50_fpn.py

@@ -0,0 +1,120 @@
+# model settings
+model = dict(
+    type='MaskRCNN',
+    pretrained='torchvision://resnet50',
+    backbone=dict(
+        type='ResNet',
+        depth=50,
+        num_stages=4,
+        out_indices=(0, 1, 2, 3),
+        frozen_stages=1,
+        norm_cfg=dict(type='BN', requires_grad=True),
+        norm_eval=True,
+        style='pytorch'),
+    neck=dict(
+        type='FPN',
+        in_channels=[256, 512, 1024, 2048],
+        out_channels=256,
+        num_outs=5),
+    rpn_head=dict(
+        type='RPNHead',
+        in_channels=256,
+        feat_channels=256,
+        anchor_generator=dict(
+            type='AnchorGenerator',
+            scales=[8],
+            ratios=[0.5, 1.0, 2.0],
+            strides=[4, 8, 16, 32, 64]),
+        bbox_coder=dict(
+            type='DeltaXYWHBBoxCoder',
+            target_means=[.0, .0, .0, .0],
+            target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(
+            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
+        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
+    roi_head=dict(
+        type='StandardRoIHead',
+        bbox_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]),
+        bbox_head=dict(
+            type='Shared2FCBBoxHead',
+            in_channels=256,
+            fc_out_channels=1024,
+            roi_feat_size=7,
+            num_classes=80,
+            bbox_coder=dict(
+                type='DeltaXYWHBBoxCoder',
+                target_means=[0., 0., 0., 0.],
+                target_stds=[0.1, 0.1, 0.2, 0.2]),
+            reg_class_agnostic=False,
+            loss_cls=dict(
+                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
+            loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
+        mask_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=14, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]),
+        mask_head=dict(
+            type='FCNMaskHead',
+            num_convs=4,
+            in_channels=256,
+            conv_out_channels=256,
+            num_classes=80,
+            loss_mask=dict(
+                type='CrossEntropyLoss', use_mask=True, loss_weight=1.0))),
+    # model training and testing settings
+    train_cfg=dict(
+        rpn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.7,
+                neg_iou_thr=0.3,
+                min_pos_iou=0.3,
+                match_low_quality=True,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=256,
+                pos_fraction=0.5,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=False),
+            allowed_border=-1,
+            pos_weight=-1,
+            debug=False),
+        rpn_proposal=dict(
+            nms_pre=2000,
+            max_per_img=1000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.5,
+                neg_iou_thr=0.5,
+                min_pos_iou=0.5,
+                match_low_quality=True,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=512,
+                pos_fraction=0.25,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=True),
+            mask_size=28,
+            pos_weight=-1,
+            debug=False)),
+    test_cfg=dict(
+        rpn=dict(
+            nms_pre=1000,
+            max_per_img=1000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.5),
+            max_per_img=100,
+            mask_thr_binary=0.5)))

configs/_base_/schedules/schedule_1x.py

@@ -0,0 +1,11 @@
+# optimizer
+optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)
+optimizer_config = dict(grad_clip=None)
+# learning policy
+lr_config = dict(
+    policy='step',
+    warmup='linear',
+    warmup_iters=500,
+    warmup_ratio=1e-6,  # 0.001
+    step=[8, 11])
+runner = dict(type='EpochBasedRunner', max_epochs=12)

configs/_base_/schedules/schedule_20e.py

@@ -0,0 +1,11 @@
+# optimizer
+optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)
+optimizer_config = dict(grad_clip=None)
+# learning policy
+lr_config = dict(
+    policy='step',
+    warmup='linear',
+    warmup_iters=500,
+    warmup_ratio=0.001,
+    step=[16, 19])
+runner = dict(type='EpochBasedRunner', max_epochs=20)

configs/_base_/schedules/schedule_2x.py

@@ -0,0 +1,11 @@
+# optimizer
+optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)
+optimizer_config = dict(grad_clip=None)
+# learning policy
+lr_config = dict(
+    policy='step',
+    warmup='linear',
+    warmup_iters=500,
+    warmup_ratio=0.001,
+    step=[16, 22])
+runner = dict(type='EpochBasedRunner', max_epochs=24)

configs/dsb/mask_rcnn_geovig_m_fpn_1x_dsb.py

@@ -0,0 +1,102 @@
+_base_ = [
+    '../_base_/models/mask_rcnn_r50_fpn.py',
+    '../_base_/schedules/schedule_1x.py',
+    '../_base_/default_runtime.py'
+]
+
+model = dict(
+    backbone=dict(
+        _delete_=True,
+        type='geovig_m_feat',
+        style='pytorch',
+        pretrained=False,
+        use_detect_adapter=True,
+        init_cfg=None
+    ),
+    neck=dict(
+        type='FPN',
+        in_channels=[64, 128, 256, 384],
+        out_channels=256,
+        num_outs=5,
+        norm_cfg=dict(type='GN', num_groups=32, requires_grad=True)
+    ),
+    roi_head=dict(
+        bbox_head=dict(num_classes=1),
+        mask_head=dict(num_classes=1)
+    ),
+    test_cfg=dict(
+        rpn=dict(nms_pre=1000, max_per_img=1000, nms=dict(type='nms', iou_threshold=0.7), min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.3),
+            max_per_img=200,
+            mask_thr_binary=0.5
+        )
+    )
+)
+
+# Dataset settings
+dataset_type = 'DSBDataset'
+data_root = '../../data/dsb2018/stage1_train'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations', with_bbox=True, with_mask=False),
+    dict(type='LoadMasksFromAnn'),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=1,
+    workers_per_gpu=4,
+    train=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file='',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file='',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file='',
+        pipeline=test_pipeline))
+
+# Fine-tuning config
+load_from = 'geovig_m_seg_4G/epoch_12.pth'
+
+evaluation = dict(metric='mAP')
+
+optimizer = dict(_delete_=True, type='AdamW', lr=0.0001, weight_decay=0.05)
+optimizer_config = dict(_delete_=True, grad_clip=dict(max_norm=1.0, norm_type=2))
+
+runner = dict(type='EpochBasedRunner', max_epochs=20)
+lr_config = dict(
+    policy='step',
+    warmup='linear',
+    warmup_iters=500,
+    warmup_ratio=0.001,
+    step=[16, 19])

configs/dsb/mask_rcnn_mobilevig_m_fpn_1x_dsb.py

@@ -0,0 +1,74 @@
+_base_ = [
+    '../mask_rcnn_mobilevig_m_fpn_1x_coco.py'
+]
+
+model = dict(
+    roi_head=dict(
+        bbox_head=dict(num_classes=1),
+        mask_head=dict(num_classes=1)
+    ),
+    test_cfg=dict(
+        rpn=dict(nms_pre=1000, max_per_img=1000, nms=dict(type='nms', iou_threshold=0.7), min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.3),
+            max_per_img=200,
+            mask_thr_binary=0.5
+        )
+    )
+)
+
+# Dataset settings
+dataset_type = 'DSBDataset'
+data_root = '../../data/dsb2018/stage1_train'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations', with_bbox=True, with_mask=False),
+    dict(type='LoadMasksFromAnn'),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=1,
+    workers_per_gpu=4,
+    train=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file='',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file='',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file='',
+        pipeline=test_pipeline))
+
+# MobileViG Pretrained Weights
+load_from = 'geovig_m_seg_4G/det_mobilevig_m_62_8.pth'
+optimizer = dict(_delete_=True, type='AdamW', lr=0.0001, weight_decay=0.05)

configs/dsb/mask_rcnn_resnet50_fpn_1x_dsb.py

@@ -0,0 +1,74 @@
+_base_ = [
+    '../_base_/models/mask_rcnn_r50_fpn.py',
+    '../_base_/schedules/schedule_1x.py',
+    '../_base_/default_runtime.py'
+]
+
+model = dict(
+    roi_head=dict(
+        bbox_head=dict(num_classes=1),
+        mask_head=dict(num_classes=1)
+    ),
+    test_cfg=dict(
+        rpn=dict(nms_pre=1000, max_per_img=1000, nms=dict(type='nms', iou_threshold=0.7), min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.3),
+            max_per_img=200,
+            mask_thr_binary=0.5
+        )
+    )
+)
+
+# Dataset settings
+dataset_type = 'DSBDataset'
+data_root = '../../data/dsb2018/stage1_train'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations', with_bbox=True, with_mask=False),
+    dict(type='LoadMasksFromAnn'),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=1,
+    workers_per_gpu=4,
+    train=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file='',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file='',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file='',
+        pipeline=test_pipeline))
+
+load_from = 'https://download.openmmlab.com/mmdetection/v2.0/mask_rcnn/mask_rcnn_r50_fpn_1x_coco/mask_rcnn_r50_fpn_1x_coco_20200205-d4b0c5d6.pth'
+optimizer = dict(type='SGD', lr=0.002, momentum=0.9, weight_decay=0.0001)

configs/dsb/mask_rcnn_swin_t_fpn_1x_dsb.py

@@ -0,0 +1,93 @@
+_base_ = [
+    '../_base_/models/mask_rcnn_r50_fpn.py',
+    '../_base_/schedules/schedule_1x.py',
+    '../_base_/default_runtime.py'
+]
+
+model = dict(
+    backbone=dict(
+        _delete_=True,
+        type='SwinTransformer',
+        embed_dims=96,
+        depths=[2, 2, 6, 2],
+        num_heads=[3, 6, 12, 24],
+        window_size=7,
+        mlp_ratio=4,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.,
+        attn_drop_rate=0.,
+        drop_path_rate=0.2,
+        patch_norm=True,
+        out_indices=(0, 1, 2, 3),
+        with_cp=False,
+        convert_weights=True,
+    ),
+    neck=dict(in_channels=[96, 192, 384, 768]),
+    roi_head=dict(
+        bbox_head=dict(num_classes=1),
+        mask_head=dict(num_classes=1)
+    ),
+    test_cfg=dict(
+        rpn=dict(nms_pre=1000, max_per_img=1000, nms=dict(type='nms', iou_threshold=0.7), min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.3),
+            max_per_img=200,
+            mask_thr_binary=0.5
+        )
+    )
+)
+
+# Dataset settings
+dataset_type = 'DSBDataset'
+data_root = '../../data/dsb2018/stage1_train'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations', with_bbox=True, with_mask=False),
+    dict(type='LoadMasksFromAnn'),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=1,
+    workers_per_gpu=4,
+    train=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file='',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file='',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file='',
+        pipeline=test_pipeline))
+
+load_from = 'https://download.openmmlab.com/mmdetection/v2.0/swin/mask_rcnn_swin-t-p4-w7_fpn_1x_coco/mask_rcnn_swin-t-p4-w7_fpn_1x_coco_20210902_120937-9d6b7cfa.pth'
+optimizer = dict(_delete_=True, type='AdamW', lr=0.0001, weight_decay=0.05)

configs/kvasir/mask_rcnn_fastvit_sa12_fpn_1x_kvasir.py

@@ -0,0 +1,82 @@
+_base_ = [
+    '../_base_/models/mask_rcnn_r50_fpn.py',
+    '../_base_/datasets/coco_instance.py',
+    '../_base_/schedules/schedule_1x.py', 
+    '../_base_/default_runtime.py'
+]
+
+# Requires importing fastvit_backbone in main/test scripts
+custom_imports = dict(imports=['fastvit_backbone'], allow_failed_imports=False)
+
+model = dict(
+    backbone=dict(
+        _delete_=True,
+        type='fastvit_sa12',
+        init_cfg=dict(type='Pretrained', checkpoint='geovig_m_5e4_8G_300_80_70/fastvit_sa12.pth.tar')
+    ),
+    neck=dict(
+        type='FPN',
+        in_channels=[64, 128, 256, 512],
+        out_channels=256,
+        num_outs=5,
+        norm_cfg=dict(type='GN', num_groups=32, requires_grad=True)
+    ),
+    roi_head=dict(
+        bbox_head=dict(num_classes=1),
+        mask_head=dict(num_classes=1)
+    )
+)
+
+dataset_type = 'KvasirDataset'
+data_root = '../../data/Kvasir-SEG'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations', with_bbox=True, with_mask=False),
+    dict(type='LoadMasksFromAnn'),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+
+data = dict(
+    samples_per_gpu=1,
+    workers_per_gpu=4,
+    train=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=test_pipeline))
+
+evaluation = dict(metric='mAP', save_best='mAP')
+optimizer = dict(_delete_=True, type='AdamW', lr=0.0001, weight_decay=0.05)

configs/kvasir/mask_rcnn_geovig_m_fpn_1x_kvasir.py

@@ -0,0 +1,85 @@
+_base_ = [
+    '../_base_/models/mask_rcnn_r50_fpn.py',
+    '../_base_/schedules/schedule_1x.py',
+    '../_base_/default_runtime.py'
+]
+
+model = dict(
+    backbone=dict(
+        _delete_=True,
+        type='geovig_m_feat',
+        style='pytorch',
+        pretrained=False,
+        use_detect_adapter=True,
+        init_cfg=None
+    ),
+    neck=dict(
+        type='FPN',
+        in_channels=[64, 128, 256, 384],
+        out_channels=256,
+        num_outs=5,
+        norm_cfg=dict(type='GN', num_groups=32, requires_grad=True)
+    ),
+    roi_head=dict(
+        bbox_head=dict(num_classes=1),
+        mask_head=dict(num_classes=1)
+    )
+)
+
+# Dataset settings
+dataset_type = 'KvasirDataset'
+data_root = '../../data/Kvasir-SEG'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations', with_bbox=True, with_mask=False),
+    dict(type='LoadMasksFromAnn'),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=1,
+    workers_per_gpu=4,
+    train=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=test_pipeline))
+
+# Fine-tuning config
+load_from = 'geovig_m_seg_4G/epoch_12.pth'
+
+evaluation = dict(metric='mAP')
+
+optimizer = dict(_delete_=True, type='AdamW', lr=0.0001, weight_decay=0.05)
+optimizer_config = dict(_delete_=True, grad_clip=dict(max_norm=1.0, norm_type=2))

configs/kvasir/mask_rcnn_mobilevig_m_fpn_1x_kvasir.py

@@ -0,0 +1,66 @@
+_base_ = [
+    '../mask_rcnn_mobilevig_m_fpn_1x_coco.py'
+]
+
+model = dict(
+    roi_head=dict(
+        bbox_head=dict(num_classes=1),
+        mask_head=dict(num_classes=1)
+    )
+)
+
+dataset_type = 'KvasirDataset'
+data_root = '../../data/Kvasir-SEG'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations', with_bbox=True, with_mask=False),
+    dict(type='LoadMasksFromAnn'),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+
+data = dict(
+    samples_per_gpu=1,
+    workers_per_gpu=4,
+    train=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=test_pipeline))
+
+evaluation = dict(metric='mAP', save_best='mAP')
+
+# MobileViG Pretrained Weights (COCO)
+load_from = 'geovig_m_seg_4G/det_mobilevig_m_62_8.pth'

configs/kvasir/mask_rcnn_resnet50_fpn_1x_kvasir.py

@@ -0,0 +1,72 @@
+_base_ = [
+    '../_base_/models/mask_rcnn_r50_fpn.py',
+    '../_base_/datasets/coco_instance.py',
+    '../_base_/schedules/schedule_1x.py', 
+    '../_base_/default_runtime.py'
+]
+
+model = dict(
+    roi_head=dict(
+        bbox_head=dict(num_classes=1),
+        mask_head=dict(num_classes=1)
+    )
+)
+
+dataset_type = 'KvasirDataset'
+data_root = '../../data/Kvasir-SEG'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations', with_bbox=True, with_mask=False),
+    dict(type='LoadMasksFromAnn'),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+
+data = dict(
+    samples_per_gpu=1,
+    workers_per_gpu=4,
+    train=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=test_pipeline))
+
+# COCO Pretrained Weights
+load_from = 'https://download.openmmlab.com/mmdetection/v2.0/mask_rcnn/mask_rcnn_r50_fpn_1x_coco/mask_rcnn_r50_fpn_1x_coco_20200205-d4b0c5d6.pth'
+
+evaluation = dict(metric='mAP', save_best='mAP')
+
+optimizer = dict(type='SGD', lr=0.002, momentum=0.9, weight_decay=0.0001)
+

configs/kvasir/mask_rcnn_swin_t_fpn_1x_kvasir.py

@@ -0,0 +1,92 @@
+_base_ = [
+    '../_base_/models/mask_rcnn_r50_fpn.py',
+    '../_base_/datasets/coco_instance.py',
+    '../_base_/schedules/schedule_1x.py', 
+    '../_base_/default_runtime.py'
+]
+
+model = dict(
+    backbone=dict(
+        _delete_=True,
+        type='SwinTransformer',
+        embed_dims=96,
+        depths=[2, 2, 6, 2],
+        num_heads=[3, 6, 12, 24],
+        window_size=7,
+        mlp_ratio=4,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.,
+        attn_drop_rate=0.,
+        drop_path_rate=0.2,
+        patch_norm=True,
+        out_indices=(0, 1, 2, 3),
+        with_cp=False,
+        convert_weights=True,
+    ),
+    neck=dict(in_channels=[96, 192, 384, 768]),
+    roi_head=dict(
+        bbox_head=dict(num_classes=1),
+        mask_head=dict(num_classes=1)
+    )
+)
+
+dataset_type = 'KvasirDataset'
+data_root = '../../data/Kvasir-SEG'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations', with_bbox=True, with_mask=False),
+    dict(type='LoadMasksFromAnn'),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+
+data = dict(
+    samples_per_gpu=1,
+    workers_per_gpu=4,
+    train=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        img_prefix=data_root,
+        ann_file=data_root + '/kavsir_bboxes.json',
+        pipeline=test_pipeline))
+
+evaluation = dict(metric='mAP', save_best='mAP')
+optimizer = dict(_delete_=True, type='AdamW', lr=0.0001, betas=(0.9, 0.999), weight_decay=0.05,
+                 paramwise_cfg=dict(custom_keys={'absolute_pos_embed': dict(decay_mult=0.),
+                                                 'relative_position_bias_table': dict(decay_mult=0.),
+                                                 'norm': dict(decay_mult=0.)}))
+
+# We use init_cfg for backbone loading, but can also set load_from for the whole detector
+load_from = 'https://download.openmmlab.com/mmdetection/v2.0/swin/mask_rcnn_swin-t-p4-w7_fpn_1x_coco/mask_rcnn_swin-t-p4-w7_fpn_1x_coco_20210902_120937-9d6b7cfa.pth'

configs/mask_rcnn_geovig_b_fpn_1x_coco.py

@@ -0,0 +1,32 @@
+
+_base_ = [
+    '_base_/models/mask_rcnn_r50_fpn.py',
+    '_base_/datasets/coco_instance.py',
+    '_base_/schedules/schedule_1x.py',
+    '_base_/default_runtime.py'
+]
+
+model = dict(
+    backbone=dict(
+        _delete_=True,
+        type='geovig_b_feat',
+        style='pytorch',
+        pretrained='../../geovig_b_5e4_8G_300_82_38/checkpoint.pth',
+        use_detect_adapter=True, # Enable new adapter
+        init_cfg=None
+    ),
+    neck=dict(
+        type='FPN',
+        in_channels=[80, 160, 320, 512],
+        out_channels=256,
+        num_outs=5,
+        norm_cfg=dict(type='GN', num_groups=32, requires_grad=True) # GroupNorm Stability
+    ))
+
+# optimizer
+optimizer = dict(_delete_=True, type='AdamW', lr=0.0001, weight_decay=0.05) # Lower LR + AdamW
+optimizer_config = dict(_delete_=True, grad_clip=dict(max_norm=1.0, norm_type=2)) # Gradient Clipping
+
+# Evaluation: Disable BN buffer broadcast (redundant for frozen BN and causes sync hangs)
+# Evaluation: DISABLED during training to prevent DDP sync timeouts. Validate offline.
+evaluation = dict(interval=100, metric=['bbox', 'segm'])

configs/mask_rcnn_geovig_m_fpn_1x_coco.py

@@ -0,0 +1,28 @@
+
+_base_ = [
+    '_base_/models/mask_rcnn_r50_fpn.py',
+    '_base_/datasets/coco_instance.py',
+    '_base_/schedules/schedule_1x.py',
+    '_base_/default_runtime.py'
+]
+
+model = dict(
+    backbone=dict(
+        _delete_=True,
+        type='geovig_m_feat',
+        style='pytorch',
+        pretrained='../../geovig_m_5e4_8G_300_80_70/checkpoint.pth',
+        use_detect_adapter=True, # Enable new adapter
+        init_cfg=None
+    ),
+    neck=dict(
+        type='FPN',
+        in_channels=[64, 128, 256, 384],
+        out_channels=256,
+        num_outs=5,
+        norm_cfg=dict(type='GN', num_groups=32, requires_grad=True) # GroupNorm Stability
+    ))
+
+# optimizer
+optimizer = dict(_delete_=True, type='AdamW', lr=0.0001, weight_decay=0.05) # Lower LR + AdamW
+optimizer_config = dict(_delete_=True, grad_clip=dict(max_norm=1.0, norm_type=2)) # Gradient Clipping

configs/mask_rcnn_mobilevig_b_fpn_1x_coco.py

@@ -0,0 +1,24 @@
+_base_ = [
+    '_base_/models/mask_rcnn_r50_fpn.py',
+    '_base_/datasets/coco_instance.py',
+    '_base_/schedules/schedule_1x.py',
+    '_base_/default_runtime.py'
+]
+model = dict(
+    backbone=dict(
+        type='mobilevig_b_feat',
+        style='pytorch',
+        init_cfg=dict(
+            type='Pretrained',
+            checkpoint='../Pretrained_Models_MobileViG/MobileViG_B_82_6.pth.tar',
+            ),
+        ),
+    neck=dict(
+        type='FPN',
+        in_channels=[42, 84, 240, 464],
+        out_channels=256,
+        num_outs=5))
+# optimizer
+optimizer = dict(_delete_=True, type='AdamW', lr=0.0002, weight_decay=0.0001)
+optimizer_config = dict(grad_clip=None)
+

configs/mask_rcnn_mobilevig_m_fpn_1x_coco.py

@@ -0,0 +1,25 @@
+_base_ = [
+    '_base_/models/mask_rcnn_r50_fpn.py',
+    '_base_/datasets/coco_instance.py',
+    '_base_/schedules/schedule_1x.py',
+    '_base_/default_runtime.py'
+]
+# optimizer
+model = dict(
+    backbone=dict(
+        type='mobilevig_m_feat',
+        style='pytorch',
+        init_cfg=dict(
+            type='Pretrained',
+            checkpoint='../../geovig_m_5e4_8G_300_80_70/MobileViG_M_80_7.pth.tar',
+        ),
+    ),
+    neck=dict(
+        type='FPN',
+        in_channels=[42, 84, 224, 400],
+        out_channels=256,
+        num_outs=5))
+# optimizer
+optimizer = dict(_delete_=True, type='AdamW', lr=0.0002, weight_decay=0.05)
+optimizer_config = dict(grad_clip=None)
+

configs/mask_rcnn_mobilevigv2_b_fpn_1x_coco.py

@@ -0,0 +1,24 @@
+_base_ = [
+    '_base_/models/mask_rcnn_r50_fpn.py',
+    '_base_/datasets/coco_instance.py',
+    '_base_/schedules/schedule_1x.py',
+    '_base_/default_runtime.py'
+]
+model = dict(
+    backbone=dict(
+        type='mobilevigv2_b_feat',
+        style='pytorch',
+        init_cfg=dict(
+            type='Pretrained',
+            checkpoint='../Results/MobileViG_V2_B_Detection.pth',
+            ),
+        ),
+    neck=dict(
+        type='FPN',
+        in_channels=[64, 128, 256, 512],
+        out_channels=256,
+        num_outs=5))
+# optimizer
+optimizer = dict(_delete_=True, type='AdamW', lr=0.0002, weight_decay=0.0001)
+optimizer_config = dict(grad_clip=None)
+

configs/mask_rcnn_mobilevigv2_m_fpn_1x_coco.py

@@ -0,0 +1,25 @@
+_base_ = [
+    '_base_/models/mask_rcnn_r50_fpn.py',
+    '_base_/datasets/coco_instance.py',
+    '_base_/schedules/schedule_1x.py',
+    '_base_/default_runtime.py'
+]
+# optimizer
+model = dict(
+    backbone=dict(
+        type='mobilevigv2_m_feat',
+        style='pytorch',
+        init_cfg=dict(
+            type='Pretrained',
+            checkpoint='../Results/MobileViG_V2_M_Detection.pth',
+        ),
+    ),
+    neck=dict(
+        type='FPN',
+        in_channels=[32, 64, 192, 384],
+        out_channels=256,
+        num_outs=5))
+# optimizer
+optimizer = dict(_delete_=True, type='AdamW', lr=0.0002, weight_decay=0.05)
+optimizer_config = dict(grad_clip=None)
+

configs/s

@@ -0,0 +1,324 @@
+""" ImageNet Training Script
+ #train.py from vig_pytorch
+
+This is intended to be a lean and easily modifiable ImageNet training script that reproduces ImageNet
+training results with some of the latest networks and training techniques. It favours canonical PyTorch
+and standard Python style over trying to be able to 'do it all.' That said, it offers quite a few speed
+and training result improvements over the usual PyTorch example scripts. Repurpose as you see fit.
+
+This script was started from an early version of the PyTorch ImageNet example
+(https://github.com/pytorch/examples/tree/master/imagenet)
+
+Hacked together by / Copyright 2020 Ross Wightman (https://github.com/rwightman)
+
+Modified from the training script used for the Vision GNN model found in: 
+    @misc{
+        https://doi.org/10.48550/arxiv.2206.00272,
+        doi = {10.48550/ARXIV.2206.00272},
+        url = {https://arxiv.org/abs/2206.00272},
+        author = {Han, Kai and Wang, Yunhe and Guo, Jianyuan and Tang, Yehui and Wu, Enhua},
+        keywords = {Computer Vision and Pattern Recognition (cs.CV), FOS: Computer and information sciences, FOS: Computer and information sciences},
+        title = {Vision GNN: An Image is Worth Graph of Nodes},
+        publisher = {arXiv},
+        year = {2022},
+        copyright = {arXiv.org perpetual, non-exclusive license}
+    }
+"""
+# Copyright (c) OpenMMLab. All rights reserved.
+import warnings
+warnings.filterwarnings('ignore')
+import argparse
+import time
+import yaml
+import os
+import logging
+from collections import OrderedDict
+from contextlib import suppress
+from datetime import datetime
+
+import copy
+import os.path as osp
+
+import torch
+import torch.nn as nn
+import torchvision.utils
+from torch.nn.parallel import DistributedDataParallel
+import timm
+
+from timm.data import ImageDataset, resolve_data_config, Mixup, FastCollateMixup, AugMixDataset #create_loader
+# from timm.models import create_model, resume_checkpoint, convert_splitbn_model
+from timm.utils import *
+from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy, JsdCrossEntropy
+from timm.optim import create_optimizer
+from timm.scheduler import create_scheduler
+from timm.utils import NativeScaler
+
+import mmcv
+from mmcv import Config, DictAction
+from mmcv.runner import get_dist_info, init_dist
+from mmcv.utils import get_git_hash
+
+from mmdet import __version__
+from mmdet.apis import init_random_seed, set_random_seed  # , train_detector
+from mmdet.datasets import build_dataset
+from mmdet.models import build_detector
+from mmdet.utils import (collect_env, get_device, get_root_logger,
+                         setup_multi_processes, update_data_root)
+
+from mmdet_custom.apis.train import train_detector
+import mmcv_custom.runner.epoch_based_runner
+import mmcv_custom.runner.optimizer
+
+import mobilevigv2_backbone
+import geovig_backbone
+
+# # 1. Disable TF32 (H100 uses this by default, but it crashes older cuBLAS)
+# torch.backends.cuda.matmul.allow_tf32 = False
+# torch.backends.cudnn.allow_tf32 = False
+
+# # 2. Force cuBLAS to use a specific, stable workspace configuration
+# # This must match the environment variable you set earlier
+# import os
+# os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
+
+# # 3. Disable the specific cuBLASLt heuristics that cause Hopper crashes
+# os.environ['DISABLE_ADDMM_CUDA_LT'] = '1'
+
+# # 4. Optional: Force deterministic math (prevents cuBLAS from choosing risky 'fast' paths)
+# torch.use_deterministic_algorithms(True, warn_only=True)
+
+parser = argparse.ArgumentParser(description='Train a detector')
+parser.add_argument('config', help='train config file path')
+parser.add_argument('--work-dir', help='the dir to save logs and models')
+parser.add_argument('--mobilevigv2_model', default='mobilevigv2_m', type=str, metavar='MODEL',
+                help='Name of model to train (default: "mobilevigv2_m"')
+parser.add_argument(
+    '--resume-from', help='the checkpoint file to resume from')
+parser.add_argument(
+    '--auto-resume',
+    action='store_true',
+    help='resume from the latest checkpoint automatically')
+parser.add_argument(
+    '--no-validate',
+    action='store_true',
+    help='whether not to evaluate the checkpoint during training')
+group_gpus = parser.add_mutually_exclusive_group()
+group_gpus.add_argument(
+    '--gpus',
+    type=int,
+    help='(Deprecated, please use --gpu-id) number of gpus to use '
+            '(only applicable to non-distributed training)')
+group_gpus.add_argument(
+    '--gpu-ids',
+    type=int,
+    nargs='+',
+    help='(Deprecated, please use --gpu-id) ids of gpus to use '
+            '(only applicable to non-distributed training)')
+group_gpus.add_argument(
+    '--gpu-id',
+    type=int,
+    default=0,
+    help='id of gpu to use '
+            '(only applicable to non-distributed training)')
+parser.add_argument('--seed', type=int, default=None, help='random seed')
+parser.add_argument(
+    '--diff-seed',
+    action='store_true',
+    help='Whether or not set different seeds for different ranks')
+parser.add_argument(
+    '--deterministic',
+    action='store_true',
+    help='whether to set deterministic options for CUDNN backend.')
+parser.add_argument(
+    '--options',
+    nargs='+',
+    action=DictAction,
+    help='override some settings in the used config, the key-value pair '
+            'in xxx=yyy format will be merged into config file (deprecate), '
+            'change to --cfg-options instead.')
+parser.add_argument(
+    '--cfg-options',
+    nargs='+',
+    action=DictAction,
+    help='override some settings in the used config, the key-value pair '
+            'in xxx=yyy format will be merged into config file. If the value to '
+            'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
+            'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
+            'Note that the quotation marks are necessary and that no white space '
+            'is allowed.')
+parser.add_argument(
+    '--launcher',
+    choices=['none', 'pytorch', 'slurm', 'mpi'],
+    default='none',
+    help='job launcher')
+parser.add_argument('--local_rank', type=int, default=0)
+parser.add_argument(
+    '--auto-scale-lr',
+    action='store_true',
+    help='enable automatically scaling LR.')
+
+args = parser.parse_args()
+print(args)
+print("\n mobilevigv2_model = " + str(args.mobilevigv2_model))
+
+def parse_args():
+    args = parser.parse_args()
+    if 'LOCAL_RANK' not in os.environ:
+        os.environ['LOCAL_RANK'] = str(args.local_rank)
+
+    if args.options and args.cfg_options:
+        raise ValueError(
+            '--options and --cfg-options cannot be both '
+            'specified, --options is deprecated in favor of --cfg-options')
+    if args.options:
+        warnings.warn('--options is deprecated in favor of --cfg-options')
+        args.cfg_options = args.options
+
+    return args
+
+
+def main():
+    args = parse_args()
+
+    cfg = Config.fromfile(args.config)
+
+    # update data root according to MMDET_DATASETS
+    update_data_root(cfg)
+
+    if args.cfg_options is not None:
+        cfg.merge_from_dict(args.cfg_options)
+
+    if args.auto_scale_lr:
+        if 'auto_scale_lr' in cfg and \
+                'enable' in cfg.auto_scale_lr and \
+                'base_batch_size' in cfg.auto_scale_lr:
+            cfg.auto_scale_lr.enable = True
+        else:
+            warnings.warn('Can not find "auto_scale_lr" or '
+                          '"auto_scale_lr.enable" or '
+                          '"auto_scale_lr.base_batch_size" in your'
+                          ' configuration file. Please update all the '
+                          'configuration files to mmdet >= 2.24.1.')
+
+    # set multi-process settings
+    setup_multi_processes(cfg)
+
+    # set cudnn_benchmark
+    if cfg.get('cudnn_benchmark', False):
+        torch.backends.cudnn.benchmark = True
+
+    # work_dir is determined in this priority: CLI > segment in file > filename
+    if args.work_dir is not None:
+        # update configs according to CLI args if args.work_dir is not None
+        cfg.work_dir = args.work_dir
+    elif cfg.get('work_dir', None) is None:
+        # use config filename as default work_dir if cfg.work_dir is None
+        cfg.work_dir = osp.join('./work_dirs',
+                                osp.splitext(osp.basename(args.config))[0])
+    if args.resume_from is not None:
+        cfg.resume_from = args.resume_from
+    cfg.auto_resume = args.auto_resume
+    if args.gpus is not None:
+        cfg.gpu_ids = range(1)
+        warnings.warn('`--gpus` is deprecated because we only support '
+                      'single GPU mode in non-distributed training. '
+                      'Use `gpus=1` now.')
+    if args.gpu_ids is not None:
+        cfg.gpu_ids = args.gpu_ids[0:1]
+        warnings.warn('`--gpu-ids` is deprecated, please use `--gpu-id`. '
+                      'Because we only support single GPU mode in '
+                      'non-distributed training. Use the first GPU '
+                      'in `gpu_ids` now.')
+    if args.gpus is None and args.gpu_ids is None:
+        cfg.gpu_ids = [args.gpu_id]
+
+    # init distributed env first, since logger depends on the dist info.
+    if args.launcher == 'none':
+        distributed = False
+        print("NOT DISTRIBUTED LAUNCH")
+    else:
+        print("DIST LAUNCH")
+        distributed = True
+        init_dist(args.launcher, **cfg.dist_params)
+        # re-set gpu_ids with distributed training mode
+        _, world_size = get_dist_info()
+        cfg.gpu_ids = range(world_size)
+        print ("GPU IDS")
+        print(cfg.gpu_ids)
+
+    # create work_dir
+    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
+    # dump config
+    cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config)))
+    # init the logger before other steps
+    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
+    log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
+    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)
+
+    # init the meta dict to record some important information such as
+    # environment info and seed, which will be logged
+    meta = dict()
+    # log env info
+    env_info_dict = collect_env()
+    env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()])
+    dash_line = '-' * 60 + '\n'
+    logger.info('Environment info:\n' + dash_line + env_info + '\n' +
+                dash_line)
+    meta['env_info'] = env_info
+    meta['config'] = cfg.pretty_text
+    # log some basic info
+    logger.info(f'Distributed training: {distributed}')
+    logger.info(f'Config:\n{cfg.pretty_text}')
+
+    cfg.device = get_device()
+    # set random seeds
+    seed = init_random_seed(args.seed, device=cfg.device)
+    seed = seed + dist.get_rank() if args.diff_seed else seed
+    logger.info(f'Set random seed to {seed}, '
+                f'deterministic: {args.deterministic}')
+    set_random_seed(seed, deterministic=args.deterministic)
+    cfg.seed = seed
+    meta['seed'] = seed
+    meta['exp_name'] = osp.basename(args.config)
+
+    model = build_detector(
+        cfg.model,
+        train_cfg=cfg.get('train_cfg'),
+        test_cfg=cfg.get('test_cfg'))
+        # drop_connect_rate=0.1,
+        # drop_path_rate=0.1,
+        # drop_block_rate=0.1)
+    model.init_weights()
+
+    print("\n\n\nCONFIG OF MODEL: \n" + str(cfg.model) + "\n\n\n")
+
+    datasets = [build_dataset(cfg.data.train)]
+    if len(cfg.workflow) == 2:
+        val_dataset = copy.deepcopy(cfg.data.val)
+        val_dataset.pipeline = cfg.data.train.pipeline
+        datasets.append(build_dataset(val_dataset))
+    if cfg.checkpoint_config is not None:
+        # save mmdet version, config file content and class names in
+        # checkpoints as meta data
+        cfg.checkpoint_config.meta = dict(
+            mmdet_version=__version__ + get_git_hash()[:7],
+            CLASSES=datasets[0].CLASSES)
+    # add an attribute for visualization convenience
+    model.CLASSES = datasets[0].CLASSES
+
+    print("\n\n\nTRAIN DETECTOR CONFIG OF MODEL: \n" + str(cfg) + "\n\n\n")
+
+    print("\n\n\TRAIN DETECTOR MODEL: \n" + str(model) + "\n\n\n")
+
+    train_detector(
+        model,
+        datasets,
+        cfg,
+        distributed=distributed,
+        validate=(not args.no_validate),
+        timestamp=timestamp,
+        meta=meta)
+
+
+if __name__ == '__main__':
+    main()

geovig.py

@@ -0,0 +1,560 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn import Sequential as Seq
+import logging
+
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+try:
+    from timm.layers import DropPath          # timm >= 0.9
+except ImportError:
+    from timm.models.layers import DropPath   # timm 0.6.x
+from timm.models import register_model
+
+'''
+@article{han2022vision,
+  title={Vision GNN: An Image is Worth Graph of Nodes},
+  author={Han, Kai and Wang, Yunhe and Guo, Jianyuan and Tang, Yehui and Wu, Enhua},
+  journal={arXiv preprint arXiv:2206.00272},
+  year={2022}
+}
+'''
+
+def _cfg(url='', **kwargs):
+    return {
+        'url': url,
+        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
+        'crop_pct': .9, 'interpolation': 'bicubic',
+        'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD, 
+        'classifier': 'head',
+        **kwargs
+    }
+
+
+default_cfgs = {
+    'mobilevig': _cfg(crop_pct=0.9, mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD)
+}
+
+    
+class Stem(nn.Module):
+    def __init__(self, input_dim, output_dim, activation=nn.GELU):
+        super(Stem, self).__init__()
+        self.stem = nn.Sequential(
+            nn.Conv2d(input_dim, output_dim // 2, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(output_dim // 2),
+            nn.GELU(),
+            nn.Conv2d(output_dim // 2, output_dim, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(output_dim),
+            nn.GELU()   
+        )
+        
+    def forward(self, x):
+        return self.stem(x)
+    
+
+class MLP(nn.Module):
+    """
+    Implementation of MLP with 1*1 convolutions.
+    Input: tensor with shape [B, C, H, W]
+    """
+
+    def __init__(self, in_features, hidden_features=None,
+                 out_features=None, drop=0., mid_conv=False):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.mid_conv = mid_conv
+        self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
+        self.act = nn.GELU()
+        self.fc2 = nn.Conv2d(hidden_features, out_features, 1)
+        self.drop = nn.Dropout(drop)
+
+        if self.mid_conv:
+            self.mid = nn.Conv2d(hidden_features, hidden_features, kernel_size=3, stride=1, padding=1,
+                                 groups=hidden_features)
+            self.mid_norm = nn.BatchNorm2d(hidden_features)
+
+        self.norm1 = nn.BatchNorm2d(hidden_features)
+        self.norm2 = nn.BatchNorm2d(out_features)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.norm1(x)
+        x = self.act(x)
+
+        if self.mid_conv:
+            x_mid = self.mid(x)
+            x_mid = self.mid_norm(x_mid)
+            x = self.act(x_mid)
+        x = self.drop(x)
+
+        x = self.fc2(x)
+        x = self.norm2(x)
+
+        x = self.drop(x)
+        return x
+
+    
+class InvertedResidual(nn.Module):
+    def __init__(self, dim, mlp_ratio=4., drop=0., drop_path=0., use_layer_scale=True, layer_scale_init_value=1e-5):
+        super().__init__()
+
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = MLP(in_features=dim, hidden_features=mlp_hidden_dim, drop=drop, mid_conv=True)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. \
+            else nn.Identity()
+        self.use_layer_scale = use_layer_scale
+        if use_layer_scale:
+            self.layer_scale_2 = nn.Parameter(
+                layer_scale_init_value * torch.ones(dim).unsqueeze(-1).unsqueeze(-1), requires_grad=True)
+
+    def forward(self, x):
+        if self.use_layer_scale:
+            x = x + self.drop_path(self.layer_scale_2 * self.mlp(x))
+        else:
+            x = x + self.drop_path(self.mlp(x))
+        return x
+    
+
+# -------------------- New graph utilities and GeoViG additions --------------------
+
+def img_to_graph(x):
+    """
+    Converts (B, C, H, W) -> (B, N, C) and generates a grid adjacency (8-neighbor-ish).
+    For efficiency this returns a single edge_index shared for the batch.
+    """
+    B, C, H, W = x.shape
+    device = x.device
+    N = H * W
+
+    # Flatten features: (B, C, H, W) -> (B, N, C)
+    x_flat = x.flatten(2).transpose(1, 2)  # (B, N, C)
+
+    # Generate grid edges (simplified 8-neighbor connectivity)
+    idx = torch.arange(N, device=device).view(H, W)
+    edges = []
+
+    # Right neighbors
+    if W > 1:
+        edges.append(torch.stack([idx[:, :-1].flatten(), idx[:, 1:].flatten()], dim=0))
+        # Left (will be mirrored later)
+    # Down neighbors
+    if H > 1:
+        edges.append(torch.stack([idx[:-1, :].flatten(), idx[1:, :].flatten()], dim=0))
+    # Diagonals
+    if H > 1 and W > 1:
+        edges.append(torch.stack([idx[:-1, :-1].flatten(), idx[1:, 1:].flatten()], dim=0))
+        edges.append(torch.stack([idx[1:, :-1].flatten(), idx[:-1, 1:].flatten()], dim=0))
+
+    if len(edges) == 0:
+        edge_index = torch.zeros((2, 0), dtype=torch.long, device=device)
+    else:
+        edge_index = torch.cat(edges, dim=1)
+        # Make undirected by mirroring
+        edge_index = torch.cat([edge_index, edge_index.flip(0)], dim=1)
+
+    return x_flat, edge_index
+
+
+class GraphMRConv(nn.Module):
+    """
+    Graph version of Max-Relative Graph Convolution.
+    Operates on (B, N, C) and edge_index (2, E).
+    Uses PyTorch's index_reduce_ with 'amax' for true max aggregation.
+    """
+    def __init__(self, in_channels, out_channels, concat=False):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.concat = concat
+        linear_in = 2 * in_channels if concat else in_channels
+        self.nn = nn.Sequential(
+            nn.Linear(linear_in, out_channels),
+            nn.GELU()
+        )
+
+    def forward(self, x, edge_index):
+        # x: (B, N, C)
+        B, N, C = x.shape
+        device = x.device
+
+        # Reshape to (B*N, C) for batched processing
+        x_flat = x.view(-1, C)
+
+        if edge_index.numel() == 0:
+            # No edges: fallback to concatenation with zeros
+            # If no edges, no aggregation, so the update is zero.
+            # We apply the non-linearity to a zero tensor to maintain shape and type.
+            return self.nn(torch.zeros_like(x_flat)).view(B, N, -1)
+
+        row, col = edge_index  # (E,), (E,)
+        num_edges = row.size(0)
+
+        # Create batched edge index: (2, B*E)
+        batch_offsets = torch.arange(B, device=device) * N
+        edge_offsets = batch_offsets.view(-1, 1).repeat(1, num_edges).view(-1)
+        
+        row_batch = row.repeat(B) + edge_offsets
+        col_batch = col.repeat(B) + edge_offsets
+
+        # Gather features
+        # Gather features
+        x_flat_c = x_flat.contiguous() # Ensure memory layout for scatter
+        # Optimization: Don't compute x_j - x_i explicitly (huge tensor).
+        # max(x_j - x_i) = max(x_j) - x_i
+        x_j = x_flat_c[col_batch]  # (B*E, C)
+
+        # Initialize aggregation buffer
+        aggr = torch.full_like(x_flat, -1e9)  # (B*N, C)
+        
+        # Vectorized Max Aggregation of Neighbors
+        index_expanded = row_batch.view(-1, 1).expand(-1, C)
+        aggr.scatter_reduce_(0, index_expanded, x_j, reduce='amax', include_self=True)
+        
+        # Replace -1e9 with 0 (although strictly if we subtract x_i later, we need to handle no-neighbor case carefully)
+        # If include_self=True, then x_i is included in x_j? No. 
+        # include_self=True in scatter_reduce_ means "reduction includes the initial value of self (aggr)".
+        # Our initial value is -1e9. So it works as a safe initializer.
+        
+        aggr = torch.where(aggr == -1e9, torch.zeros_like(aggr), aggr)
+
+        if self.concat:
+            rel = aggr - x_flat
+            out = torch.cat([x_flat, rel], dim=-1)  # (B*N, 2C)
+        else:
+            out = aggr - x_flat                     # (B*N, C)
+
+        return self.nn(out).view(B, N, -1)
+
+
+class GraphMobileViGBlock(nn.Module):
+    """
+    Graph MobileViG block: GraphMRConv + FFN for (B,N,C) representations.
+    """
+    def __init__(self, dim, drop=0.0, concat=False):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.graph_conv = GraphMRConv(dim, dim, concat=concat)
+        self.norm2 = nn.LayerNorm(dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * 4),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim),
+            nn.Dropout(drop)
+        )
+
+    def forward(self, x, edge_index):
+        # x: (B, N, C)
+        shortcut = x
+        x = self.norm1(x)
+        x = self.graph_conv(x, edge_index)
+        x = x + shortcut
+
+        shortcut = x
+        x = self.norm2(x)
+        x = self.ffn(x)
+        x = x + shortcut
+        return x
+
+
+class SpreadEdgePool(nn.Module):
+    """
+    Geometry-Aware SpreadEdgePool: uses edge scores (diffusion distance proxy) as 
+    attention weights to guide feature pooling. Ensures nodes with high structural 
+    diversity contribute more to the pooled representation.
+    """
+    def __init__(self, in_channels, ratio=0.5):
+        super().__init__()
+        self.ratio = ratio
+        self.in_channels = in_channels
+
+    def forward(self, x, edge_index):
+        """
+        x: (B, N, C)
+        edge_index: (2, E)
+        Returns:
+            x_pooled: (B, N_new, C) with geometry-aware weighting
+            new_edge_index: (2, E_new)
+        """
+        B, N, C = x.shape
+        device = x.device
+        num_keep = max(1, int(N * self.ratio))
+
+        # Fallback if no edges
+        if edge_index.numel() == 0:
+            x_pooled = F.adaptive_avg_pool1d(x.transpose(1, 2), num_keep).transpose(1, 2)
+            idx = torch.arange(num_keep, device=device)
+            new_edge_index = torch.stack([idx[:-1], idx[1:]], dim=0) if num_keep > 1 else torch.zeros((2, 0), dtype=torch.long, device=device)
+            return x_pooled, new_edge_index
+
+        row, col = edge_index  # (E,), (E,)
+
+        # Geometry-Aware Scoring: Diffusion distance proxy via feature space
+        # Optimization: Calculate distance WITHOUT materializing (B, E, C) difference tensor.
+        # |A - B|^2 = |A|^2 + |B|^2 - 2(A . B)
+        
+        # Precompute norms: (B, N, 1)
+        x_sq = (x ** 2).sum(dim=-1, keepdim=True)
+        
+        # Gather norms: (B, E, 1)
+        sq_i = x_sq[:, row, :]
+        sq_j = x_sq[:, col, :]
+        
+        # Gather features just for dot (this still uses memory but dot reduces it immediately)
+        x_row = x[:, row, :]  # (B, E, C)
+        x_col = x[:, col, :]  # (B, E, C)
+        dot = (x_row * x_col).sum(dim=-1, keepdim=True) # (B, E, 1)
+        
+        # Combine
+        dist_sq = sq_i + sq_j - 2 * dot
+        dist = torch.sqrt(F.relu(dist_sq) + 1e-6).squeeze(-1) # (B, E)
+
+        # Spread score: high score = high distance = high structural diversity
+        # Original: "similarity = torch.exp(-dist)", "edge_scores = 1.0 / (similarity + 1e-6)"
+        # This means edge_scores = exp(dist). So higher distance means higher score.
+        # Using distance directly achieves this.
+        edge_scores = dist # Linear distance directly. 
+        
+        # Map edge scores to node importance: sum scores of incident edges
+        # Average over batch for global topology understanding
+        avg_edge_scores = edge_scores.mean(dim=0)  # (E,)
+        
+        # Aggregate edge importance to nodes via index_add_
+        node_importance = torch.zeros(N, device=device)
+        node_importance.index_add_(0, row, avg_edge_scores)  # Sum incoming edge scores per node
+        
+        # Normalize importance scores to [0, 1] via sigmoid
+        node_weights = torch.sigmoid(node_importance).view(1, N, 1)  # (1, N, 1) for broadcasting
+
+        # Geometry-Aware Pooling: Weight features by their geometric importance before aggregation
+        x_weighted = x * node_weights  # (B, N, C), broadcast multiply
+        
+        # Adaptive pooling now prioritizes geometrically important features
+        x_pooled = F.adaptive_avg_pool1d(x_weighted.transpose(1, 2), num_keep).transpose(1, 2)  # (B, num_keep, C)
+
+        # Build heuristic new adjacency: simple chain / local ring
+        idx = torch.arange(num_keep, device=device)
+        if num_keep > 1:
+            left = idx[:-1]
+            right = idx[1:]
+            new_edge_index = torch.cat([torch.stack([left, right], dim=0), torch.stack([right, left], dim=0)], dim=1)
+        else:
+            new_edge_index = torch.zeros((2, 0), dtype=torch.long, device=device)
+
+        return x_pooled, new_edge_index
+
+
+class Downsample(nn.Module):
+    """
+    Convolution-based downsample
+    """
+    def __init__(self, in_dim, out_dim):
+        super().__init__()        
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_dim, out_dim, 3, stride=2, padding=1),
+            nn.BatchNorm2d(out_dim),
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+    
+class GeoViG(nn.Module):
+    """
+    Hybrid Grid->Graph GeoViG model. Can be used for classification or as a backbone for detection.
+    """
+    def __init__(self,
+                 local_blocks=[2, 2],
+                 local_channels=[42, 84],
+                 graph_blocks=[2],
+                 graph_channels=[168],
+                 num_classes=1000,
+                 pool_ratio=0.5,
+                 distillation=False,
+                 features_only=False,
+                 out_indices=None,
+                 pretrained=None,
+                 use_concat=False,
+                 **kwargs):
+        super().__init__()
+
+        self.features_only = features_only
+        self.out_indices = out_indices
+        self.distillation = distillation
+        self.pretrained = pretrained
+        self.use_concat = use_concat
+
+        # Stage 0: stem
+        self.stem = Stem(input_dim=3, output_dim=local_channels[0])
+
+        # Grid stages
+        self.grid_stages = nn.ModuleList()
+        in_c = local_channels[0]
+        
+        self.feature_info = [dict(num_chs=in_c, reduction=2, module=f'stem')]
+
+        for i, out_c in enumerate(local_channels):
+            layers = []
+            if i > 0:
+                layers.append(Downsample(in_c, out_c))
+                in_c = out_c
+            for _ in range(local_blocks[i]):
+                layers.append(InvertedResidual(dim=out_c, mlp_ratio=4, drop_path=0.0))
+            self.grid_stages.append(nn.Sequential(*layers))
+            self.feature_info.append(dict(num_chs=out_c, reduction=2**(i+2), module=f'grid_stages.{i}'))
+
+        # Transition: grid -> graph
+        self.to_graph_pool = SpreadEdgePool(in_channels=in_c, ratio=pool_ratio)
+        self.graph_dim = in_c
+
+        # Graph stages
+        self.graph_proj = nn.Linear(self.graph_dim, graph_channels[0]) if graph_channels else nn.Identity()
+        self.graph_dim = graph_channels[0] if graph_channels else self.graph_dim
+
+        self.graph_stages = nn.ModuleList()
+        for _ in range(graph_blocks[0] if len(graph_blocks) > 0 else 0):
+            self.graph_stages.append(GraphMobileViGBlock(dim=self.graph_dim, concat=use_concat))
+        
+        self.feature_info.append(dict(num_chs=self.graph_dim, reduction=2**(len(local_channels)+1), module='graph_stages'))
+
+        if not self.features_only:
+            # Head for classification
+            self.norm = nn.LayerNorm(self.graph_dim)
+            self.head = nn.Linear(self.graph_dim, num_classes)
+            if self.distillation:
+                self.head_dist = nn.Linear(self.graph_dim, num_classes)
+        else:
+            # For torchvision FPN
+            self._out_channels = []
+            if self.out_indices:
+                for i in self.out_indices:
+                    # The indices for feature_info are shifted by 1 compared to grid_stages
+                    self._out_channels.append(self.feature_info[i+1]['num_chs'])
+
+        if self.pretrained:
+            self._load_pretrained_weights()
+
+    @property
+    def out_channels(self):
+        if not self.features_only:
+            raise ValueError("out_channels is only available in features_only mode.")
+        return self._out_channels
+
+    def _load_pretrained_weights(self):
+        """Load pretrained weights from checkpoint."""
+        if self.pretrained and isinstance(self.pretrained, str):
+            logging.info(f"Loading pretrained weights from {self.pretrained}")
+            state_dict = torch.load(self.pretrained, map_location='cpu', weights_only=False)
+            
+            if 'state_dict' in state_dict:
+                state_dict = state_dict['state_dict']
+            elif 'model' in state_dict:
+                state_dict = state_dict['model']
+
+            model_dict = self.state_dict()
+            
+            # Filter out keys that are not in the current model, have different shapes, or belong to the classification head
+            pretrained_dict = {
+                k: v for k, v in state_dict.items() 
+                if k in model_dict and v.shape == model_dict[k].shape and 'head' not in k
+            }
+            
+            model_dict.update(pretrained_dict) 
+            self.load_state_dict(model_dict, strict=False)
+            logging.info(f"Loaded {len(pretrained_dict)} keys from pretrained weights.")
+
+    def train(self, mode=True):
+        super().train(mode)
+        # For detection, keep BatchNorm in eval mode
+        if self.features_only:
+            for m in self.modules():
+                if isinstance(m, nn.BatchNorm2d):
+                    m.eval()
+
+    def forward(self, x):
+        outs = {}
+        x = self.stem(x)
+
+        for i, stage in enumerate(self.grid_stages):
+            x = stage(x)
+            if self.features_only and self.out_indices and i in self.out_indices:
+                outs[str(i)] = x
+        
+        # Graph processing
+        B, C, H, W = x.shape
+        x_nodes, edge_index = img_to_graph(x)
+        x_nodes, edge_index = self.to_graph_pool(x_nodes, edge_index)
+        x_nodes = self.graph_proj(x_nodes)
+        
+        for block in self.graph_stages:
+            x_nodes = block(x_nodes, edge_index)
+        
+        if self.features_only:
+            graph_stage_idx = len(self.grid_stages)
+            if self.out_indices and graph_stage_idx in self.out_indices:
+                num_nodes = x_nodes.shape[1]
+                new_h = int(num_nodes**0.5)
+                new_w = int(num_nodes**0.5)
+                if new_h * new_w != num_nodes:
+                    new_h = new_w = int(num_nodes**0.5)
+                    x_nodes = x_nodes[:, :new_h*new_w, :]
+                
+                x_graph = x_nodes.transpose(1, 2).view(B, self.graph_dim, new_h, new_w)
+                outs[str(graph_stage_idx)] = x_graph
+            
+            return outs
+
+        # Global readout (mean) for classification
+        x_global = x_nodes.mean(dim=1)
+        x_out = self.norm(x_global)
+        
+        if self.distillation:
+            x_main = self.head(x_out)
+            x_dist = self.head_dist(x_out)
+            if self.training:
+                return x_main, x_dist
+            else:
+                return (x_main + x_dist) / 2
+        else:
+            return self.head(x_out)
+
+
+# Register a small GeoViG variant for timm-style usage
+@register_model
+def geovig_ti(pretrained=False, **kwargs):
+    model = GeoViG(local_blocks=[3, 3, 3], local_channels=[42, 84, 168],
+                   graph_blocks=[3], graph_channels=[256],
+                   pool_ratio=0.5, use_concat=True,
+                   pretrained=pretrained, **kwargs)
+    model.default_cfg = default_cfgs['mobilevig']
+    return model
+
+@register_model
+def geovig_s(pretrained=False, **kwargs):
+    model = GeoViG(local_blocks=[3, 3, 4], local_channels=[56, 112, 224],
+                   graph_blocks=[3], graph_channels=[288],
+                   pool_ratio=0.5, use_concat=False,
+                   pretrained=pretrained, **kwargs)
+    model.default_cfg = default_cfgs['mobilevig']
+    return model
+
+@register_model
+def geovig_m(pretrained=False, **kwargs):
+    # Target: ~10M Params (MobileViG-M: 14M)
+    model = GeoViG(local_blocks=[3, 3, 9], local_channels=[64, 128, 256],
+                   graph_blocks=[3], graph_channels=[384],
+                   pool_ratio=0.5,
+                   pretrained=pretrained, **kwargs)
+    model.default_cfg = default_cfgs['mobilevig']
+    return model
+
+@register_model
+def geovig_b(pretrained=False, **kwargs):
+    # Target: ~19M Params (MobileViG-B: 26.7M)
+    model = GeoViG(local_blocks=[5, 5, 9], local_channels=[80, 160, 320],
+                   graph_blocks=[4], graph_channels=[512],
+                   pool_ratio=0.5,
+                   pretrained=pretrained, **kwargs)
+    model.default_cfg = default_cfgs['mobilevig']
+    return model

geovig_backbone.py

@@ -0,0 +1,667 @@
+import os
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn import Sequential as Seq
+import logging
+import copy
+
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from timm.models.layers import DropPath
+from timm.models.registry import register_model
+
+try:
+    from mmdet.models.builder import BACKBONES as det_BACKBONES
+    from mmdet.utils import get_root_logger
+    from mmcv.runner import _load_checkpoint
+    has_mmdet = True
+except ImportError:
+    print("If for detection, please install mmdetection first")
+    has_mmdet = False
+
+def _cfg(url='', **kwargs):
+    return {
+        'url': url,
+        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
+        'crop_pct': .9, 'interpolation': 'bicubic',
+        'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD, 
+        'classifier': 'head',
+        **kwargs
+    }
+
+default_cfgs = {
+    'mobilevig': _cfg(crop_pct=0.9, mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD)
+}
+
+import logging
+
+logging.basicConfig(
+    level=logging.WARNING,
+    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s"
+)
+
+logger = logging.getLogger(__name__)
+    
+logger = logging.getLogger(__name__)
+
+class Stem(nn.Module):
+    def __init__(self, input_dim, output_dim, activation=nn.GELU):
+        super(Stem, self).__init__()
+        self.stem = nn.Sequential(
+            nn.Conv2d(input_dim, output_dim // 2, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(output_dim // 2),
+            nn.GELU(),
+            nn.Conv2d(output_dim // 2, output_dim, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(output_dim),
+            nn.GELU()   
+        )
+        
+    def forward(self, x):
+        return self.stem(x)
+
+class GraphAdapter(nn.Module):
+    """
+    Graph-Aware Adapter: Lightweight 3x3 DW Conv + GroupNorm + 1x1 Conv + GroupNorm.
+    Aligns graph features with grid-based detection heads.
+    """
+    def __init__(self, channels):
+        super().__init__()
+        self.adapter = nn.Sequential(
+            # Depthwise 3x3
+            nn.Conv2d(channels, channels, kernel_size=3, padding=1, groups=channels, bias=False),
+            # GroupNorm (num_groups=1 is equivalent to LayerNorm but works with 4D tensors)
+            nn.GroupNorm(1, channels), 
+            # Pointwise 1x1
+            nn.Conv2d(channels, channels, kernel_size=1, bias=False),
+            nn.GroupNorm(32, channels), # GroupNorm for stability
+            nn.GELU()
+        )
+
+    def forward(self, x):
+        return self.adapter(x)
+
+    
+
+class MLP(nn.Module):
+    """
+    Implementation of MLP with 1*1 convolutions.
+    Input: tensor with shape [B, C, H, W]
+    """
+
+    def __init__(self, in_features, hidden_features=None,
+                 out_features=None, drop=0., mid_conv=False):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.mid_conv = mid_conv
+        self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
+        self.act = nn.GELU()
+        self.fc2 = nn.Conv2d(hidden_features, out_features, 1)
+        self.drop = nn.Dropout(drop)
+
+        if self.mid_conv:
+            self.mid = nn.Conv2d(hidden_features, hidden_features, kernel_size=3, stride=1, padding=1,
+                                 groups=hidden_features)
+            self.mid_norm = nn.BatchNorm2d(hidden_features)
+
+        self.norm1 = nn.BatchNorm2d(hidden_features)
+        self.norm2 = nn.BatchNorm2d(out_features)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.norm1(x)
+        x = self.act(x)
+
+        if self.mid_conv:
+            x_mid = self.mid(x)
+            x_mid = self.mid_norm(x_mid)
+            x = self.act(x_mid)
+        x = self.drop(x)
+
+        x = self.fc2(x)
+        x = self.norm2(x)
+
+        x = self.drop(x)
+        return x
+
+    
+class InvertedResidual(nn.Module):
+    def __init__(self, dim, mlp_ratio=4., drop=0., drop_path=0., use_layer_scale=True, layer_scale_init_value=1e-5):
+        super().__init__()
+
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = MLP(in_features=dim, hidden_features=mlp_hidden_dim, drop=drop, mid_conv=True)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. \
+            else nn.Identity()
+        self.use_layer_scale = use_layer_scale
+        if use_layer_scale:
+            self.layer_scale_2 = nn.Parameter(
+                layer_scale_init_value * torch.ones(dim).unsqueeze(-1).unsqueeze(-1), requires_grad=True)
+
+    def forward(self, x):
+        if self.use_layer_scale:
+            x = x + self.drop_path(self.layer_scale_2 * self.mlp(x))
+        else:
+            x = x + self.drop_path(self.mlp(x))
+        return x
+    
+
+def img_to_graph(x):
+    """
+    Converts (B, C, H, W) -> (B, N, C) and generates a grid adjacency (8-neighbor-ish).
+    For efficiency this returns a single edge_index shared for the batch.
+    """
+    B, C, H, W = x.shape
+    device = x.device
+    N = H * W
+
+    # Flatten features: (B, C, H, W) -> (B, N, C)
+    x_flat = x.flatten(2).transpose(1, 2)  # (B, N, C)
+
+    # Generate grid edges (simplified 8-neighbor connectivity)
+    idx = torch.arange(N, device=device).view(H, W)
+    edges = []
+
+    # Right neighbors
+    if W > 1:
+        edges.append(torch.stack([idx[:, :-1].flatten(), idx[:, 1:].flatten()], dim=0))
+        # Left (will be mirrored later)
+    # Down neighbors
+    if H > 1:
+        edges.append(torch.stack([idx[:-1, :].flatten(), idx[1:, :].flatten()], dim=0))
+    # Diagonals
+    if H > 1 and W > 1:
+        edges.append(torch.stack([idx[:-1, :-1].flatten(), idx[1:, 1:].flatten()], dim=0))
+        edges.append(torch.stack([idx[1:, :-1].flatten(), idx[:-1, 1:].flatten()], dim=0))
+
+    if len(edges) == 0:
+        edge_index = torch.zeros((2, 0), dtype=torch.long, device=device)
+    else:
+        edge_index = torch.cat(edges, dim=1)
+        # Make undirected by mirroring
+        edge_index = torch.cat([edge_index, edge_index.flip(0)], dim=1)
+
+    return x_flat, edge_index
+
+
+class GraphMRConv(nn.Module):
+    """
+    Graph version of Max-Relative Graph Convolution.
+    Operates on (B, N, C) and edge_index (2, E).
+    Uses PyTorch's index_reduce_ with 'amax' for true max aggregation.
+    """
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        # OPTIMIZATION: Removed concatenation (2*C). Now just Linear(C -> C).
+        # This reduces parameters significantly and speeds up the projection.
+        self.nn = nn.Sequential(
+            nn.Linear(in_channels, out_channels),
+            nn.GELU()
+        )
+
+    def forward(self, x, edge_index):
+        # x: (B, N, C)
+        B, N, C = x.shape
+        device = x.device
+
+        # Reshape to (B*N, C) for batched processing
+        x_flat = x.view(-1, C)
+
+        if edge_index.numel() == 0:
+            # No edges: fallback to identity-like behavior with zeros
+            # Return processed x_flat reshaped, or zeros if safer
+            # Ideally graph conv with no edges is just local projection of 0-vector (since it's rel diff)
+            # max(x_j) would be empty -> -inf -> replaced by 0
+            # then -x_i
+            # so effectively -x_i passed through nn
+            return self.nn(-x).view(B, N, -1)
+
+        row, col = edge_index  # (E,), (E,)
+        num_edges = row.size(0)
+
+        # Create batched edge index: (2, B*E)
+        batch_offsets = torch.arange(B, device=device) * N
+        edge_offsets = batch_offsets.view(-1, 1).repeat(1, num_edges).view(-1)
+        
+        row_batch = row.repeat(B) + edge_offsets
+        col_batch = col.repeat(B) + edge_offsets
+
+        # Gather features
+        x_flat_c = x_flat.contiguous() 
+        # Optimization: Don't compute x_j - x_i explicitly (huge tensor).
+        # max(x_j - x_i) = max(x_j) - x_i
+        x_j = x_flat_c[col_batch]  # (B*E, C)
+
+        # Initialize aggregation buffer
+        aggr = torch.full_like(x_flat, -1e9)  # (B*N, C)
+        
+        # Vectorized Max Aggregation of Neighbors
+        index_expanded = row_batch.view(-1, 1).expand(-1, C)
+        aggr.scatter_reduce_(0, index_expanded, x_j, reduce='amax', include_self=True)
+        
+        # Replace -1e9 with 0 
+        aggr = torch.where(aggr == -1e9, torch.zeros_like(aggr), aggr)
+        
+        # Now apply the "- x_i" part of "max(x_j) - x_i"
+        aggr = aggr - x_flat
+
+        # Output only the aggregated update
+        return self.nn(aggr).view(B, N, -1)
+
+
+class GraphMobileViGBlock(nn.Module):
+    """
+    Graph MobileViG block: GraphMRConv + FFN for (B,N,C) representations.
+    """
+    def __init__(self, dim, drop=0.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.graph_conv = GraphMRConv(dim, dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * 4),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim),
+            nn.Dropout(drop)
+        )
+
+    def forward(self, x, edge_index):
+        # x: (B, N, C)
+        shortcut = x
+        x = self.norm1(x)
+        x = self.graph_conv(x, edge_index)
+        x = x + shortcut
+
+        shortcut = x
+        x = self.norm2(x)
+        x = self.ffn(x)
+        x = x + shortcut
+        return x
+
+
+class SpreadEdgePool(nn.Module):
+    """
+    Geometry-Aware SpreadEdgePool
+    """
+    def __init__(self, in_channels, ratio=0.5):
+        super().__init__()
+        self.ratio = ratio
+        self.in_channels = in_channels
+
+    def forward(self, x, edge_index, shape=None):
+        B, N, C = x.shape
+        device = x.device
+
+        # 2D Scaling Path (For Detection) - STABLE MODE
+        # Skips distance-based weighting to avoid instability in detection fine-tuning
+        if shape is not None:
+            H, W = shape
+            if H * W == N:
+                x_2d = x.transpose(1, 2).view(B, C, H, W)
+                
+                # Force Stride 2 for Detection (Target Stride 32 from Stride 16 input)
+                # ratio=0.25 implies 1/4 area -> stride 2
+                stride = 2 if self.ratio <= 0.25 else 1 
+                
+                x_pooled_2d = F.avg_pool2d(x_2d, kernel_size=stride, stride=stride)
+                
+                # Regenerate Grid Graph
+                x_pooled, new_edge_index = img_to_graph(x_pooled_2d)
+                return x_pooled, new_edge_index
+
+        # Default 1D Graph Pooling (original classification logic)
+        num_keep = max(1, int(N * self.ratio))
+        
+        # Only compute weights if strictly needed (1D path)
+        if edge_index.numel() > 0:
+            row, col = edge_index
+            x_sq = (x ** 2).sum(dim=-1, keepdim=True)
+            sq_i = x_sq[:, row, :]
+            sq_j = x_sq[:, col, :]
+            x_row = x[:, row, :]
+            x_col = x[:, col, :]
+            dot = (x_row * x_col).sum(dim=-1, keepdim=True)
+            dist_sq = sq_i + sq_j - 2 * dot
+            dist = torch.sqrt(F.relu(dist_sq) + 1e-6).squeeze(-1)
+            avg_edge_scores = dist.mean(dim=0)
+            node_importance = torch.zeros(N, device=device)
+            node_importance.index_add_(0, row, avg_edge_scores)
+            node_weights = torch.sigmoid(node_importance).view(1, N, 1)
+        else:
+            node_weights = torch.ones((1, N, 1), device=device)
+
+        x_weighted = x * node_weights
+        num_keep = max(1, int(N * self.ratio))
+        
+        x_pooled = F.adaptive_avg_pool1d(x_weighted.transpose(1, 2), num_keep).transpose(1, 2)
+
+        if edge_index.numel() == 0:
+            idx = torch.arange(num_keep, device=device)
+            new_edge_index = torch.stack([idx[:-1], idx[1:]], dim=0) if num_keep > 1 else torch.zeros((2, 0), dtype=torch.long, device=device)
+            return x_pooled, new_edge_index
+
+        idx = torch.arange(num_keep, device=device)
+        if num_keep > 1:
+            left = idx[:-1]
+            right = idx[1:]
+            new_edge_index = torch.cat([torch.stack([left, right], dim=0), torch.stack([right, left], dim=0)], dim=1)
+        else:
+            new_edge_index = torch.zeros((2, 0), dtype=torch.long, device=device)
+
+        return x_pooled, new_edge_index
+
+
+class Downsample(nn.Module):
+    """
+    Convolution-based downsample
+    """
+    def __init__(self, in_dim, out_dim):
+        super().__init__()        
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_dim, out_dim, 3, stride=2, padding=1),
+            nn.BatchNorm2d(out_dim),
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+    
+class GeoViG(nn.Module):
+    """
+    Hybrid Grid->Graph GeoViG model. Can be used for classification or as a backbone for detection.
+    """
+    def __init__(self,
+                 local_blocks=[2, 2],
+                 local_channels=[42, 84],
+                 graph_blocks=[2],
+                 graph_channels=[168],
+                 num_classes=1000,
+                 pool_ratio=0.5,
+                 distillation=False,
+                 features_only=False,
+                 out_indices=None,
+                 pretrained=None,
+                 use_detect_adapter=False,
+                 **kwargs):
+        super().__init__()
+
+        self.features_only = features_only
+        self.pool_ratio = pool_ratio
+        self.out_indices = out_indices
+        self.distillation = distillation
+        self.pretrained = pretrained
+        self.use_detect_adapter = use_detect_adapter
+
+        # Stage 0: stem
+        self.stem = Stem(input_dim=3, output_dim=local_channels[0])
+
+        # Grid stages
+        self.grid_stages = nn.ModuleList()
+        in_c = local_channels[0]
+        
+        self.feature_info = [dict(num_chs=in_c, reduction=2, module=f'stem')]
+
+        for i, out_c in enumerate(local_channels):
+            layers = []
+            if i > 0:
+                layers.append(Downsample(in_c, out_c))
+                in_c = out_c
+            for _ in range(local_blocks[i]):
+                layers.append(InvertedResidual(dim=out_c, mlp_ratio=4, drop_path=0.0))
+            self.grid_stages.append(nn.Sequential(*layers))
+            self.feature_info.append(dict(num_chs=out_c, reduction=2**(i+2), module=f'grid_stages.{i}'))
+
+        # Transition: grid -> graph
+        self.to_graph_pool = SpreadEdgePool(in_channels=in_c, ratio=pool_ratio)
+        self.graph_dim = in_c
+
+        # Graph stages
+        self.graph_proj = nn.Linear(self.graph_dim, graph_channels[0]) if graph_channels else nn.Identity()
+        self.graph_dim = graph_channels[0] if graph_channels else self.graph_dim
+
+        self.graph_stages = nn.ModuleList()
+        for _ in range(graph_blocks[0] if len(graph_blocks) > 0 else 0):
+            self.graph_stages.append(GraphMobileViGBlock(dim=self.graph_dim))
+        
+        self.feature_info.append(dict(num_chs=self.graph_dim, reduction=2**(len(local_channels)+1), module='graph_stages'))
+
+        if not self.features_only:
+            # Head for classification
+            self.norm = nn.LayerNorm(self.graph_dim)
+            self.head = nn.Linear(self.graph_dim, num_classes)
+            if self.distillation:
+                self.head_dist = nn.Linear(self.graph_dim, num_classes)
+        else:
+            # For torchvision FPN
+            self._out_channels = []
+            if self.out_indices:
+                for i in self.out_indices:
+                    # The indices for feature_info are shifted by 1 compared to grid_stages
+                    # features_info: [stem(idx0), grid0(idx1), grid1(idx2), grid2(idx3), graph(idx4)]
+                    # grid_stages: [grid0, grid1, grid2] (indices 0, 1, 2)
+                    pass
+
+        # Detection Adapter Initialization
+        if self.use_detect_adapter and self.features_only:
+            # Determine channels automatically
+            # Stem: local_channels[0] (but usually we tap grid stages)
+            # Grid Stages output: local_channels
+            # Graph Stage output: graph_channels[0]
+            
+            # Map indices to channels
+            # We assume out_indices=[0, 1, 2, 3] corresponds to [grid0, grid1, grid2, graph]
+            # Verify this mapping convention from forward loop
+            self.detect_adapters = nn.ModuleList()
+            
+            # This is a bit dynamic because out_indices can vary, but assuming standard GeoViG-M/B
+            # and standard out_indices=[0, 1, 2, 3]
+            
+            current_channels = []
+            # Indices refer to stages
+            # 0 -> grid_stages[0] (local_channels[0])
+            # 1 -> grid_stages[1] (local_channels[1])
+            # 2 -> grid_stages[2] (local_channels[2])
+            # 3 -> graph_stage (graph_channels[0]) (if present)
+            
+            all_channels = local_channels[:]
+            if graph_channels:
+                all_channels.append(graph_channels[0])
+                
+            for idx in (self.out_indices or []):
+                if idx < len(all_channels):
+                    self.detect_adapters.append(GraphAdapter(all_channels[idx]))
+                else:
+                    # Fallback or error
+                    pass
+
+        self.init_weights()
+        if pretrained:
+            self._load_pretrained_weights()
+        
+        # self = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self)
+
+    def init_weights(self):
+        if self.pretrained and isinstance(self.pretrained, str):
+            logger = get_root_logger()
+            print(f"Pretrained weights being being loaded from {self.pretrained}")
+            logger.warn(f'Pretrained weights being loaded from {self.pretrained}')
+            self._load_pretrained_weights()
+        else:
+            # Init random weights if needed or rely on default init
+            pass
+
+    def _load_pretrained_weights(self):
+        """Load pretrained weights from checkpoint."""
+        if self.pretrained and isinstance(self.pretrained, str):
+            logging.info(f"Loading pretrained weights from {self.pretrained}")
+            try:
+                # Use mmcv loader if available
+                ckpt = _load_checkpoint(self.pretrained, logger=logging.getLogger(), map_location='cpu')
+            except:
+                ckpt = torch.load(self.pretrained, map_location='cpu')
+
+            print("ckpt keys: ", ckpt.keys())
+            if 'state_dict_ema' in ckpt:
+                state_dict = ckpt['state_dict_ema']
+            elif 'state_dict' in ckpt:
+                state_dict = ckpt['state_dict']
+            elif 'model' in ckpt:
+                state_dict = ckpt['model']
+            else:
+                state_dict = ckpt
+
+            model_dict = self.state_dict()
+            
+            # Simple permissive loading
+            pretrained_dict = {}
+            for k, v in state_dict.items():
+                if k in model_dict:
+                    if v.shape == model_dict[k].shape:
+                        # Exclude classification head logic
+                        if 'head' not in k:
+                            pretrained_dict[k] = v
+                    else:
+                        print(f"Shape mismatch for {k}: {v.shape} vs {model_dict[k].shape}")
+
+            missing, unexpected = self.load_state_dict(pretrained_dict, strict=False)
+            print("missing_keys: ", missing)
+            print("unexpected_keys: ", unexpected)
+
+    @torch.no_grad()
+    def train(self, mode=True):
+        super().train(mode)
+        # For detection, keep BatchNorm in eval mode
+        for m in self.modules():
+            if isinstance(m, nn.BatchNorm2d):
+                m.eval()
+
+    def forward(self, x):
+        outs = []
+        x = self.stem(x)
+
+        for i, stage in enumerate(self.grid_stages):
+            x = stage(x)
+            if self.features_only and self.out_indices and i in self.out_indices:
+                outs.append(x)
+        
+        # Graph processing
+        B, C, H, W = x.shape
+        x_nodes, edge_index = img_to_graph(x)
+        # We need to know the stride used by SpreadEdgePool to reconstruct the shape later
+        # Logic matches SpreadEdgePool: stride 2 if ratio <= 0.25 else 1
+        pool_stride = 2 if self.pool_ratio <= 0.25 else 1
+        
+        x_nodes, edge_index = self.to_graph_pool(x_nodes, edge_index, shape=(H, W))
+        x_nodes = self.graph_proj(x_nodes)
+        
+        for block in self.graph_stages:
+            x_nodes = block(x_nodes, edge_index)
+        
+        if self.features_only:
+            graph_stage_idx = len(self.grid_stages)
+            if self.out_indices and graph_stage_idx in self.out_indices:
+                # Calculate correct new dimensions
+                new_h = H // pool_stride
+                new_w = W // pool_stride
+                
+                # Verify shape matches
+                if new_h * new_w != x_nodes.shape[1]:
+                    logger.warning(
+                        "Shape mismatch in Graph Stage reshape. Expected %dx%d=%d, got %d",
+                        new_h, new_w, new_h * new_w, x_nodes.shape[1]
+                    )
+                    # Fallback to square if totally broken (should not happen with stable pooling)
+                    num_nodes = x_nodes.shape[1]
+                    new_h = int(num_nodes**0.5)
+                    new_w = int(num_nodes**0.5)
+                    if new_h * new_w != num_nodes:
+                         x_nodes = x_nodes[:, :new_h*new_w, :]
+
+                x_graph = x_nodes.transpose(1, 2).view(B, self.graph_dim, new_h, new_w)
+                outs.append(x_graph)
+            
+            # Apply Adapters and Clamping
+            final_outs = []
+            for i, feat in enumerate(outs):
+                # Stability Fix: Clamp
+                feat = torch.clamp(feat, min=-10.0, max=10.0)
+                
+                # Adaptation: Apply GraphAdapter if enabled
+                if self.use_detect_adapter and hasattr(self, 'detect_adapters') and i < len(self.detect_adapters):
+                     feat = self.detect_adapters[i](feat)
+                
+                final_outs.append(feat)
+            
+            return tuple(final_outs)
+
+        # Global readout (mean) for classification
+        x_global = x_nodes.mean(dim=1)
+        x_out = self.norm(x_global)
+        
+        if self.distillation:
+            x_main = self.head(x_out)
+            x_dist = self.head_dist(x_out)
+            if self.training:
+                return x_main, x_dist
+            else:
+                return (x_main + x_dist) / 2
+        else:
+            return self.head(x_out)
+
+
+if has_mmdet:
+    @det_BACKBONES.register_module()
+    def geovig_m_feat(pretrained=False, **kwargs):
+        if pretrained == 'torchvision://resnet50':
+            pretrained = True
+
+        # Absolute path relative to this file — works from any CWD
+        _here = os.path.dirname(os.path.abspath(__file__))
+        real_pretrained = os.path.normpath(
+            os.path.join(_here, '../final_model_ipa_pth/pth/geovig_m_5e4_8G_300_80_70/checkpoint.pth')
+        )
+        # Skip backbone pre-loading if path missing; init_detector loads the full ckpt anyway
+        if pretrained is True:
+            load_path = real_pretrained if os.path.exists(real_pretrained) else None
+        elif isinstance(pretrained, str) and os.path.exists(pretrained):
+            load_path = pretrained
+        else:
+            load_path = None
+
+        model = GeoViG(local_blocks=[3, 3, 9], local_channels=[64, 128, 256],
+                   graph_blocks=[3], graph_channels=[384],
+                   pool_ratio=0.25,
+                   features_only=True,
+                   out_indices=[0, 1, 2, 3],
+                   pretrained=load_path)
+        model.default_cfg = default_cfgs['mobilevig']
+        return model
+
+    @det_BACKBONES.register_module()
+    def geovig_b_feat(pretrained=False, **kwargs):
+        if pretrained == 'torchvision://resnet50':
+            pretrained = True
+
+        _here = os.path.dirname(os.path.abspath(__file__))
+        real_pretrained = os.path.normpath(
+            os.path.join(_here, '../final_model_ipa_pth/pth/geovig_b_5e4_8G_300_82_38/checkpoint.pth')
+        )
+        if pretrained is True:
+            load_path = real_pretrained if os.path.exists(real_pretrained) else None
+        elif isinstance(pretrained, str) and os.path.exists(pretrained):
+            load_path = pretrained
+        else:
+            load_path = None
+
+        model = GeoViG(local_blocks=[5, 5, 9], local_channels=[80, 160, 320],
+                   graph_blocks=[4], graph_channels=[512],
+                   pool_ratio=0.25,
+                   features_only=True,
+                   out_indices=[0, 1, 2, 3],
+                   pretrained=load_path)
+        model.default_cfg = default_cfgs['mobilevig']
+        return model
+

geovig_det_backbone.py

@@ -0,0 +1,676 @@
+import os
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn import Sequential as Seq
+import logging
+import copy
+
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from timm.models.layers import DropPath
+from timm.models.registry import register_model
+
+try:
+    from mmdet.models.builder import BACKBONES as det_BACKBONES
+    from mmdet.utils import get_root_logger
+    from mmcv.runner import _load_checkpoint
+    has_mmdet = True
+    
+    # Register dummy medical datasets for configuration compatibility
+    from mmdet.datasets.builder import DATASETS
+    from mmdet.datasets.custom import CustomDataset
+    
+    @DATASETS.register_module(force=True)
+    class KvasirDataset(CustomDataset):
+        CLASSES = ('polyp',)
+        
+    @DATASETS.register_module(force=True)
+    class DSBDataset(CustomDataset):
+        CLASSES = ('nucleus',)
+except ImportError:
+    print("If for detection, please install mmdetection first")
+    has_mmdet = False
+
+def _cfg(url='', **kwargs):
+    return {
+        'url': url,
+        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
+        'crop_pct': .9, 'interpolation': 'bicubic',
+        'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD, 
+        'classifier': 'head',
+        **kwargs
+    }
+
+default_cfgs = {
+    'mobilevig': _cfg(crop_pct=0.9, mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD)
+}
+
+import logging
+
+logging.basicConfig(
+    level=logging.WARNING,
+    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s"
+)
+
+logger = logging.getLogger(__name__)
+    
+logger = logging.getLogger(__name__)
+
+class Stem(nn.Module):
+    def __init__(self, input_dim, output_dim, activation=nn.GELU):
+        super(Stem, self).__init__()
+        self.stem = nn.Sequential(
+            nn.Conv2d(input_dim, output_dim // 2, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(output_dim // 2),
+            nn.GELU(),
+            nn.Conv2d(output_dim // 2, output_dim, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(output_dim),
+            nn.GELU()   
+        )
+        
+    def forward(self, x):
+        return self.stem(x)
+
+class GraphAdapter(nn.Module):
+    """
+    Graph-Aware Adapter: Lightweight 3x3 DW Conv + GroupNorm + 1x1 Conv + GroupNorm.
+    Aligns graph features with grid-based detection heads.
+    """
+    def __init__(self, channels):
+        super().__init__()
+        self.adapter = nn.Sequential(
+            # Depthwise 3x3
+            nn.Conv2d(channels, channels, kernel_size=3, padding=1, groups=channels, bias=False),
+            # GroupNorm (num_groups=1 is equivalent to LayerNorm but works with 4D tensors)
+            nn.GroupNorm(1, channels), 
+            # Pointwise 1x1
+            nn.Conv2d(channels, channels, kernel_size=1, bias=False),
+            nn.GroupNorm(32, channels), # GroupNorm for stability
+            nn.GELU()
+        )
+
+    def forward(self, x):
+        return self.adapter(x)
+
+    
+
+class MLP(nn.Module):
+    """
+    Implementation of MLP with 1*1 convolutions.
+    Input: tensor with shape [B, C, H, W]
+    """
+
+    def __init__(self, in_features, hidden_features=None,
+                 out_features=None, drop=0., mid_conv=False):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.mid_conv = mid_conv
+        self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
+        self.act = nn.GELU()
+        self.fc2 = nn.Conv2d(hidden_features, out_features, 1)
+        self.drop = nn.Dropout(drop)
+
+        if self.mid_conv:
+            self.mid = nn.Conv2d(hidden_features, hidden_features, kernel_size=3, stride=1, padding=1,
+                                 groups=hidden_features)
+            self.mid_norm = nn.BatchNorm2d(hidden_features)
+
+        self.norm1 = nn.BatchNorm2d(hidden_features)
+        self.norm2 = nn.BatchNorm2d(out_features)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.norm1(x)
+        x = self.act(x)
+
+        if self.mid_conv:
+            x_mid = self.mid(x)
+            x_mid = self.mid_norm(x_mid)
+            x = self.act(x_mid)
+        x = self.drop(x)
+
+        x = self.fc2(x)
+        x = self.norm2(x)
+
+        x = self.drop(x)
+        return x
+
+    
+class InvertedResidual(nn.Module):
+    def __init__(self, dim, mlp_ratio=4., drop=0., drop_path=0., use_layer_scale=True, layer_scale_init_value=1e-5):
+        super().__init__()
+
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = MLP(in_features=dim, hidden_features=mlp_hidden_dim, drop=drop, mid_conv=True)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. \
+            else nn.Identity()
+        self.use_layer_scale = use_layer_scale
+        if use_layer_scale:
+            self.layer_scale_2 = nn.Parameter(
+                layer_scale_init_value * torch.ones(dim).unsqueeze(-1).unsqueeze(-1), requires_grad=True)
+
+    def forward(self, x):
+        if self.use_layer_scale:
+            x = x + self.drop_path(self.layer_scale_2 * self.mlp(x))
+        else:
+            x = x + self.drop_path(self.mlp(x))
+        return x
+    
+
+def img_to_graph(x):
+    """
+    Converts (B, C, H, W) -> (B, N, C) and generates a grid adjacency (8-neighbor-ish).
+    For efficiency this returns a single edge_index shared for the batch.
+    """
+    B, C, H, W = x.shape
+    device = x.device
+    N = H * W
+
+    # Flatten features: (B, C, H, W) -> (B, N, C)
+    x_flat = x.flatten(2).transpose(1, 2)  # (B, N, C)
+
+    # Generate grid edges (simplified 8-neighbor connectivity)
+    idx = torch.arange(N, device=device).view(H, W)
+    edges = []
+
+    # Right neighbors
+    if W > 1:
+        edges.append(torch.stack([idx[:, :-1].flatten(), idx[:, 1:].flatten()], dim=0))
+        # Left (will be mirrored later)
+    # Down neighbors
+    if H > 1:
+        edges.append(torch.stack([idx[:-1, :].flatten(), idx[1:, :].flatten()], dim=0))
+    # Diagonals
+    if H > 1 and W > 1:
+        edges.append(torch.stack([idx[:-1, :-1].flatten(), idx[1:, 1:].flatten()], dim=0))
+        edges.append(torch.stack([idx[1:, :-1].flatten(), idx[:-1, 1:].flatten()], dim=0))
+
+    if len(edges) == 0:
+        edge_index = torch.zeros((2, 0), dtype=torch.long, device=device)
+    else:
+        edge_index = torch.cat(edges, dim=1)
+        # Make undirected by mirroring
+        edge_index = torch.cat([edge_index, edge_index.flip(0)], dim=1)
+
+    return x_flat, edge_index
+
+
+class GraphMRConv(nn.Module):
+    """
+    Graph version of Max-Relative Graph Convolution.
+    Operates on (B, N, C) and edge_index (2, E).
+    Uses PyTorch's index_reduce_ with 'amax' for true max aggregation.
+    """
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        # OPTIMIZATION: Removed concatenation (2*C). Now just Linear(C -> C).
+        # This reduces parameters significantly and speeds up the projection.
+        self.nn = nn.Sequential(
+            nn.Linear(in_channels, out_channels),
+            nn.GELU()
+        )
+
+    def forward(self, x, edge_index):
+        # x: (B, N, C)
+        B, N, C = x.shape
+        device = x.device
+
+        # Reshape to (B*N, C) for batched processing
+        x_flat = x.view(-1, C)
+
+        if edge_index.numel() == 0:
+            # No edges: fallback to identity-like behavior with zeros
+            # Return processed x_flat reshaped, or zeros if safer
+            # Ideally graph conv with no edges is just local projection of 0-vector (since it's rel diff)
+            # max(x_j) would be empty -> -inf -> replaced by 0
+            # then -x_i
+            # so effectively -x_i passed through nn
+            return self.nn(-x).view(B, N, -1)
+
+        row, col = edge_index  # (E,), (E,)
+        num_edges = row.size(0)
+
+        # Create batched edge index: (2, B*E)
+        batch_offsets = torch.arange(B, device=device) * N
+        edge_offsets = batch_offsets.view(-1, 1).repeat(1, num_edges).view(-1)
+        
+        row_batch = row.repeat(B) + edge_offsets
+        col_batch = col.repeat(B) + edge_offsets
+
+        # Gather features
+        x_flat_c = x_flat.contiguous() 
+        # Optimization: Don't compute x_j - x_i explicitly (huge tensor).
+        # max(x_j - x_i) = max(x_j) - x_i
+        x_j = x_flat_c[col_batch]  # (B*E, C)
+
+        # Initialize aggregation buffer
+        aggr = torch.full_like(x_flat, -1e9)  # (B*N, C)
+        
+        # Vectorized Max Aggregation of Neighbors
+        index_expanded = row_batch.view(-1, 1).expand(-1, C)
+        aggr.scatter_reduce_(0, index_expanded, x_j, reduce='amax', include_self=True)
+        
+        # Replace -1e9 with 0 
+        aggr = torch.where(aggr == -1e9, torch.zeros_like(aggr), aggr)
+        
+        # Now apply the "- x_i" part of "max(x_j) - x_i"
+        aggr = aggr - x_flat
+
+        # Output only the aggregated update
+        return self.nn(aggr).view(B, N, -1)
+
+
+class GraphMobileViGBlock(nn.Module):
+    """
+    Graph MobileViG block: GraphMRConv + FFN for (B,N,C) representations.
+    """
+    def __init__(self, dim, drop=0.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.graph_conv = GraphMRConv(dim, dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * 4),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim),
+            nn.Dropout(drop)
+        )
+
+    def forward(self, x, edge_index):
+        # x: (B, N, C)
+        shortcut = x
+        x = self.norm1(x)
+        x = self.graph_conv(x, edge_index)
+        x = x + shortcut
+
+        shortcut = x
+        x = self.norm2(x)
+        x = self.ffn(x)
+        x = x + shortcut
+        return x
+
+
+class SpreadEdgePool(nn.Module):
+    """
+    Geometry-Aware SpreadEdgePool
+    """
+    def __init__(self, in_channels, ratio=0.5):
+        super().__init__()
+        self.ratio = ratio
+        self.in_channels = in_channels
+
+    def forward(self, x, edge_index, shape=None):
+        B, N, C = x.shape
+        device = x.device
+
+        # 2D Scaling Path (For Detection) - STABLE MODE
+        # Skips distance-based weighting to avoid instability in detection fine-tuning
+        if shape is not None:
+            H, W = shape
+            if H * W == N:
+                x_2d = x.transpose(1, 2).view(B, C, H, W)
+                
+                # Force Stride 2 for Detection (Target Stride 32 from Stride 16 input)
+                # ratio=0.25 implies 1/4 area -> stride 2
+                stride = 2 if self.ratio <= 0.25 else 1 
+                
+                x_pooled_2d = F.avg_pool2d(x_2d, kernel_size=stride, stride=stride)
+                
+                # Regenerate Grid Graph
+                x_pooled, new_edge_index = img_to_graph(x_pooled_2d)
+                return x_pooled, new_edge_index
+
+        # Default 1D Graph Pooling (original classification logic)
+        num_keep = max(1, int(N * self.ratio))
+        
+        # Only compute weights if strictly needed (1D path)
+        if edge_index.numel() > 0:
+            row, col = edge_index
+            x_sq = (x ** 2).sum(dim=-1, keepdim=True)
+            sq_i = x_sq[:, row, :]
+            sq_j = x_sq[:, col, :]
+            x_row = x[:, row, :]
+            x_col = x[:, col, :]
+            dot = (x_row * x_col).sum(dim=-1, keepdim=True)
+            dist_sq = sq_i + sq_j - 2 * dot
+            dist = torch.sqrt(F.relu(dist_sq) + 1e-6).squeeze(-1)
+            avg_edge_scores = dist.mean(dim=0)
+            node_importance = torch.zeros(N, device=device)
+            node_importance.index_add_(0, row, avg_edge_scores)
+            node_weights = torch.sigmoid(node_importance).view(1, N, 1)
+        else:
+            node_weights = torch.ones((1, N, 1), device=device)
+
+        x_weighted = x * node_weights
+        num_keep = max(1, int(N * self.ratio))
+        
+        x_pooled = F.adaptive_avg_pool1d(x_weighted.transpose(1, 2), num_keep).transpose(1, 2)
+
+        if edge_index.numel() == 0:
+            idx = torch.arange(num_keep, device=device)
+            new_edge_index = torch.stack([idx[:-1], idx[1:]], dim=0) if num_keep > 1 else torch.zeros((2, 0), dtype=torch.long, device=device)
+            return x_pooled, new_edge_index
+
+        idx = torch.arange(num_keep, device=device)
+        if num_keep > 1:
+            left = idx[:-1]
+            right = idx[1:]
+            new_edge_index = torch.cat([torch.stack([left, right], dim=0), torch.stack([right, left], dim=0)], dim=1)
+        else:
+            new_edge_index = torch.zeros((2, 0), dtype=torch.long, device=device)
+
+        return x_pooled, new_edge_index
+
+
+class Downsample(nn.Module):
+    """
+    Convolution-based downsample
+    """
+    def __init__(self, in_dim, out_dim):
+        super().__init__()        
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_dim, out_dim, 3, stride=2, padding=1),
+            nn.BatchNorm2d(out_dim),
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+    
+class GeoViG(nn.Module):
+    """
+    Hybrid Grid->Graph GeoViG model. Can be used for classification or as a backbone for detection.
+    """
+    def __init__(self,
+                 local_blocks=[2, 2],
+                 local_channels=[42, 84],
+                 graph_blocks=[2],
+                 graph_channels=[168],
+                 num_classes=1000,
+                 pool_ratio=0.5,
+                 distillation=False,
+                 features_only=False,
+                 out_indices=None,
+                 pretrained=None,
+                 use_detect_adapter=False,
+                 **kwargs):
+        super().__init__()
+
+        self.features_only = features_only
+        self.pool_ratio = pool_ratio
+        self.out_indices = out_indices
+        self.distillation = distillation
+        self.pretrained = pretrained
+        self.use_detect_adapter = use_detect_adapter
+
+        # Stage 0: stem
+        self.stem = Stem(input_dim=3, output_dim=local_channels[0])
+
+        # Grid stages
+        self.grid_stages = nn.ModuleList()
+        in_c = local_channels[0]
+        
+        self.feature_info = [dict(num_chs=in_c, reduction=2, module=f'stem')]
+
+        for i, out_c in enumerate(local_channels):
+            layers = []
+            if i > 0:
+                layers.append(Downsample(in_c, out_c))
+                in_c = out_c
+            for _ in range(local_blocks[i]):
+                layers.append(InvertedResidual(dim=out_c, mlp_ratio=4, drop_path=0.0))
+            self.grid_stages.append(nn.Sequential(*layers))
+            self.feature_info.append(dict(num_chs=out_c, reduction=2**(i+2), module=f'grid_stages.{i}'))
+
+        # Transition: grid -> graph
+        self.to_graph_pool = SpreadEdgePool(in_channels=in_c, ratio=pool_ratio)
+        self.graph_dim = in_c
+
+        # Graph stages
+        self.graph_proj = nn.Linear(self.graph_dim, graph_channels[0]) if graph_channels else nn.Identity()
+        self.graph_dim = graph_channels[0] if graph_channels else self.graph_dim
+
+        self.graph_stages = nn.ModuleList()
+        for _ in range(graph_blocks[0] if len(graph_blocks) > 0 else 0):
+            self.graph_stages.append(GraphMobileViGBlock(dim=self.graph_dim))
+        
+        self.feature_info.append(dict(num_chs=self.graph_dim, reduction=2**(len(local_channels)+1), module='graph_stages'))
+
+        if not self.features_only:
+            # Head for classification
+            self.norm = nn.LayerNorm(self.graph_dim)
+            self.head = nn.Linear(self.graph_dim, num_classes)
+            if self.distillation:
+                self.head_dist = nn.Linear(self.graph_dim, num_classes)
+        else:
+            # For torchvision FPN
+            self._out_channels = []
+            if self.out_indices:
+                for i in self.out_indices:
+                    # The indices for feature_info are shifted by 1 compared to grid_stages
+                    # features_info: [stem(idx0), grid0(idx1), grid1(idx2), grid2(idx3), graph(idx4)]
+                    # grid_stages: [grid0, grid1, grid2] (indices 0, 1, 2)
+                    pass
+
+        # Detection Adapter Initialization
+        if self.use_detect_adapter and self.features_only:
+            # Determine channels automatically
+            # Stem: local_channels[0] (but usually we tap grid stages)
+            # Grid Stages output: local_channels
+            # Graph Stage output: graph_channels[0]
+            
+            # Map indices to channels
+            # We assume out_indices=[0, 1, 2, 3] corresponds to [grid0, grid1, grid2, graph]
+            # Verify this mapping convention from forward loop
+            self.detect_adapters = nn.ModuleList()
+            
+            # This is a bit dynamic because out_indices can vary, but assuming standard GeoViG-M/B
+            # and standard out_indices=[0, 1, 2, 3]
+            
+            current_channels = []
+            # Indices refer to stages
+            # 0 -> grid_stages[0] (local_channels[0])
+            # 1 -> grid_stages[1] (local_channels[1])
+            # 2 -> grid_stages[2] (local_channels[2])
+            # 3 -> graph_stage (graph_channels[0]) (if present)
+            
+            all_channels = local_channels[:]
+            if graph_channels:
+                all_channels.append(graph_channels[0])
+                
+            for idx in (self.out_indices or []):
+                if idx < len(all_channels):
+                    self.detect_adapters.append(GraphAdapter(all_channels[idx]))
+                else:
+                    # Fallback or error
+                    pass
+
+        self.init_weights()
+        if pretrained:
+            self._load_pretrained_weights()
+        
+        # self = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self)
+
+    def init_weights(self):
+        if self.pretrained and isinstance(self.pretrained, str):
+            logger = get_root_logger()
+            print(f"Pretrained weights being being loaded from {self.pretrained}")
+            logger.warn(f'Pretrained weights being loaded from {self.pretrained}')
+            self._load_pretrained_weights()
+        else:
+            # Init random weights if needed or rely on default init
+            pass
+
+    def _load_pretrained_weights(self):
+        """Load pretrained weights from checkpoint."""
+        if self.pretrained and isinstance(self.pretrained, str):
+            logging.info(f"Loading pretrained weights from {self.pretrained}")
+            try:
+                # Use mmcv loader if available
+                ckpt = _load_checkpoint(self.pretrained, logger=logging.getLogger(), map_location='cpu')
+            except:
+                ckpt = torch.load(self.pretrained, map_location='cpu')
+
+            print("ckpt keys: ", ckpt.keys())
+            if 'state_dict_ema' in ckpt:
+                state_dict = ckpt['state_dict_ema']
+            elif 'state_dict' in ckpt:
+                state_dict = ckpt['state_dict']
+            elif 'model' in ckpt:
+                state_dict = ckpt['model']
+            else:
+                state_dict = ckpt
+
+            model_dict = self.state_dict()
+            
+            # Simple permissive loading
+            pretrained_dict = {}
+            for k, v in state_dict.items():
+                if k in model_dict:
+                    if v.shape == model_dict[k].shape:
+                        # Exclude classification head logic
+                        if 'head' not in k:
+                            pretrained_dict[k] = v
+                    else:
+                        print(f"Shape mismatch for {k}: {v.shape} vs {model_dict[k].shape}")
+
+            missing, unexpected = self.load_state_dict(pretrained_dict, strict=False)
+            print("missing_keys: ", missing)
+            print("unexpected_keys: ", unexpected)
+
+    @torch.no_grad()
+    def train(self, mode=True):
+        super().train(mode)
+        # For detection, keep BatchNorm in eval mode
+        for m in self.modules():
+            if isinstance(m, nn.BatchNorm2d):
+                m.eval()
+
+    def forward(self, x):
+        outs = []
+        x = self.stem(x)
+
+        for i, stage in enumerate(self.grid_stages):
+            x = stage(x)
+            if self.features_only and self.out_indices and i in self.out_indices:
+                outs.append(x)
+        
+        # Graph processing
+        B, C, H, W = x.shape
+        x_nodes, edge_index = img_to_graph(x)
+        # We need to know the stride used by SpreadEdgePool to reconstruct the shape later
+        # Logic matches SpreadEdgePool: stride 2 if ratio <= 0.25 else 1
+        pool_stride = 2 if self.pool_ratio <= 0.25 else 1
+        
+        x_nodes, edge_index = self.to_graph_pool(x_nodes, edge_index, shape=(H, W))
+        x_nodes = self.graph_proj(x_nodes)
+        
+        for block in self.graph_stages:
+            x_nodes = block(x_nodes, edge_index)
+        
+        if self.features_only:
+            graph_stage_idx = len(self.grid_stages)
+            if self.out_indices and graph_stage_idx in self.out_indices:
+                # Calculate correct new dimensions
+                new_h = H // pool_stride
+                new_w = W // pool_stride
+                
+                # Verify shape matches
+                if new_h * new_w != x_nodes.shape[1]:
+                    logger.warning(
+                        "Shape mismatch in Graph Stage reshape. Expected %dx%d=%d, got %d",
+                        new_h, new_w, new_h * new_w, x_nodes.shape[1]
+                    )
+                    # Fallback to square if totally broken (should not happen with stable pooling)
+                    num_nodes = x_nodes.shape[1]
+                    new_h = int(num_nodes**0.5)
+                    new_w = int(num_nodes**0.5)
+                    if new_h * new_w != num_nodes:
+                         x_nodes = x_nodes[:, :new_h*new_w, :]
+
+                x_graph = x_nodes.transpose(1, 2).view(B, self.graph_dim, new_h, new_w)
+                outs.append(x_graph)
+            
+            # Apply Adapters
+            final_outs = []
+            for i, feat in enumerate(outs):
+                # Adaptation: Apply GraphAdapter if enabled
+                if self.use_detect_adapter and hasattr(self, 'detect_adapters') and i < len(self.detect_adapters):
+                     feat = self.detect_adapters[i](feat)
+                
+                final_outs.append(feat)
+            
+            return tuple(final_outs)
+
+        # Global readout (mean) for classification
+        x_global = x_nodes.mean(dim=1)
+        x_out = self.norm(x_global)
+        
+        if self.distillation:
+            x_main = self.head(x_out)
+            x_dist = self.head_dist(x_out)
+            if self.training:
+                return x_main, x_dist
+            else:
+                return (x_main + x_dist) / 2
+        else:
+            return self.head(x_out)
+
+
+if has_mmdet:
+    @det_BACKBONES.register_module()
+    def geovig_m_feat(pretrained=False, **kwargs):
+        if pretrained == 'torchvision://resnet50':
+            pretrained = True
+
+        # Absolute path relative to this file — works from any CWD
+        _here = os.path.dirname(os.path.abspath(__file__))
+        real_pretrained = os.path.normpath(
+            os.path.join(_here, '../final_model_ipa_pth/pth/geovig_m_5e4_8G_300_80_70/checkpoint.pth')
+        )
+        # Skip backbone pre-loading if path missing; init_detector loads the full ckpt anyway
+        if pretrained is True:
+            load_path = real_pretrained if os.path.exists(real_pretrained) else None
+        elif isinstance(pretrained, str) and os.path.exists(pretrained):
+            load_path = pretrained
+        else:
+            load_path = None
+
+        model = GeoViG(local_blocks=[3, 3, 9], local_channels=[64, 128, 256],
+                   graph_blocks=[3], graph_channels=[384],
+                   pool_ratio=0.25,
+                   features_only=True,
+                   out_indices=[0, 1, 2, 3],
+                   pretrained=load_path)
+        model.default_cfg = default_cfgs['mobilevig']
+        return model
+
+    @det_BACKBONES.register_module()
+    def geovig_b_feat(pretrained=False, **kwargs):
+        if pretrained == 'torchvision://resnet50':
+            pretrained = True
+
+        _here = os.path.dirname(os.path.abspath(__file__))
+        real_pretrained = os.path.normpath(
+            os.path.join(_here, '../final_model_ipa_pth/pth/geovig_b_5e4_8G_300_82_38/checkpoint.pth')
+        )
+        if pretrained is True:
+            load_path = real_pretrained if os.path.exists(real_pretrained) else None
+        elif isinstance(pretrained, str) and os.path.exists(pretrained):
+            load_path = pretrained
+        else:
+            load_path = None
+
+        model = GeoViG(local_blocks=[5, 5, 9], local_channels=[80, 160, 320],
+                   graph_blocks=[4], graph_channels=[512],
+                   pool_ratio=0.25,
+                   features_only=True,
+                   out_indices=[0, 1, 2, 3],
+                   pretrained=load_path)
+        model.default_cfg = default_cfgs['mobilevig']
+        return model
+

imagenet_labels.json

@@ -0,0 +1,1000 @@
+["tench",
+"goldfish",
+"great white shark",
+"tiger shark",
+"hammerhead shark",
+"electric ray",
+"stingray",
+"cock",
+"hen",
+"ostrich",
+"brambling",
+"goldfinch",
+"house finch",
+"junco",
+"indigo bunting",
+"American robin",
+"bulbul",
+"jay",
+"magpie",
+"chickadee",
+"American dipper",
+"kite",
+"bald eagle",
+"vulture",
+"great grey owl",
+"fire salamander",
+"smooth newt",
+"newt",
+"spotted salamander",
+"axolotl",
+"American bullfrog",
+"tree frog",
+"tailed frog",
+"loggerhead sea turtle",
+"leatherback sea turtle",
+"mud turtle",
+"terrapin",
+"box turtle",
+"banded gecko",
+"green iguana",
+"Carolina anole",
+"desert grassland whiptail lizard",
+"agama",
+"frilled-necked lizard",
+"alligator lizard",
+"Gila monster",
+"European green lizard",
+"chameleon",
+"Komodo dragon",
+"Nile crocodile",
+"American alligator",
+"triceratops",
+"worm snake",
+"ring-necked snake",
+"eastern hog-nosed snake",
+"smooth green snake",
+"kingsnake",
+"garter snake",
+"water snake",
+"vine snake",
+"night snake",
+"boa constrictor",
+"African rock python",
+"Indian cobra",
+"green mamba",
+"sea snake",
+"Saharan horned viper",
+"eastern diamondback rattlesnake",
+"sidewinder",
+"trilobite",
+"harvestman",
+"scorpion",
+"yellow garden spider",
+"barn spider",
+"European garden spider",
+"southern black widow",
+"tarantula",
+"wolf spider",
+"tick",
+"centipede",
+"black grouse",
+"ptarmigan",
+"ruffed grouse",
+"prairie grouse",
+"peacock",
+"quail",
+"partridge",
+"grey parrot",
+"macaw",
+"sulphur-crested cockatoo",
+"lorikeet",
+"coucal",
+"bee eater",
+"hornbill",
+"hummingbird",
+"jacamar",
+"toucan",
+"duck",
+"red-breasted merganser",
+"goose",
+"black swan",
+"tusker",
+"echidna",
+"platypus",
+"wallaby",
+"koala",
+"wombat",
+"jellyfish",
+"sea anemone",
+"brain coral",
+"flatworm",
+"nematode",
+"conch",
+"snail",
+"slug",
+"sea slug",
+"chiton",
+"chambered nautilus",
+"Dungeness crab",
+"rock crab",
+"fiddler crab",
+"red king crab",
+"American lobster",
+"spiny lobster",
+"crayfish",
+"hermit crab",
+"isopod",
+"white stork",
+"black stork",
+"spoonbill",
+"flamingo",
+"little blue heron",
+"great egret",
+"bittern",
+"crane (bird)",
+"limpkin",
+"common gallinule",
+"American coot",
+"bustard",
+"ruddy turnstone",
+"dunlin",
+"common redshank",
+"dowitcher",
+"oystercatcher",
+"pelican",
+"king penguin",
+"albatross",
+"grey whale",
+"killer whale",
+"dugong",
+"sea lion",
+"Chihuahua",
+"Japanese Chin",
+"Maltese",
+"Pekingese",
+"Shih Tzu",
+"King Charles Spaniel",
+"Papillon",
+"toy terrier",
+"Rhodesian Ridgeback",
+"Afghan Hound",
+"Basset Hound",
+"Beagle",
+"Bloodhound",
+"Bluetick Coonhound",
+"Black and Tan Coonhound",
+"Treeing Walker Coonhound",
+"English foxhound",
+"Redbone Coonhound",
+"borzoi",
+"Irish Wolfhound",
+"Italian Greyhound",
+"Whippet",
+"Ibizan Hound",
+"Norwegian Elkhound",
+"Otterhound",
+"Saluki",
+"Scottish Deerhound",
+"Weimaraner",
+"Staffordshire Bull Terrier",
+"American Staffordshire Terrier",
+"Bedlington Terrier",
+"Border Terrier",
+"Kerry Blue Terrier",
+"Irish Terrier",
+"Norfolk Terrier",
+"Norwich Terrier",
+"Yorkshire Terrier",
+"Wire Fox Terrier",
+"Lakeland Terrier",
+"Sealyham Terrier",
+"Airedale Terrier",
+"Cairn Terrier",
+"Australian Terrier",
+"Dandie Dinmont Terrier",
+"Boston Terrier",
+"Miniature Schnauzer",
+"Giant Schnauzer",
+"Standard Schnauzer",
+"Scottish Terrier",
+"Tibetan Terrier",
+"Australian Silky Terrier",
+"Soft-coated Wheaten Terrier",
+"West Highland White Terrier",
+"Lhasa Apso",
+"Flat-Coated Retriever",
+"Curly-coated Retriever",
+"Golden Retriever",
+"Labrador Retriever",
+"Chesapeake Bay Retriever",
+"German Shorthaired Pointer",
+"Vizsla",
+"English Setter",
+"Irish Setter",
+"Gordon Setter",
+"Brittany Spaniel",
+"Clumber Spaniel",
+"English Springer Spaniel",
+"Welsh Springer Spaniel",
+"Cocker Spaniels",
+"Sussex Spaniel",
+"Irish Water Spaniel",
+"Kuvasz",
+"Schipperke",
+"Groenendael",
+"Malinois",
+"Briard",
+"Australian Kelpie",
+"Komondor",
+"Old English Sheepdog",
+"Shetland Sheepdog",
+"collie",
+"Border Collie",
+"Bouvier des Flandres",
+"Rottweiler",
+"German Shepherd Dog",
+"Dobermann",
+"Miniature Pinscher",
+"Greater Swiss Mountain Dog",
+"Bernese Mountain Dog",
+"Appenzeller Sennenhund",
+"Entlebucher Sennenhund",
+"Boxer",
+"Bullmastiff",
+"Tibetan Mastiff",
+"French Bulldog",
+"Great Dane",
+"St. Bernard",
+"husky",
+"Alaskan Malamute",
+"Siberian Husky",
+"Dalmatian",
+"Affenpinscher",
+"Basenji",
+"pug",
+"Leonberger",
+"Newfoundland",
+"Pyrenean Mountain Dog",
+"Samoyed",
+"Pomeranian",
+"Chow Chow",
+"Keeshond",
+"Griffon Bruxellois",
+"Pembroke Welsh Corgi",
+"Cardigan Welsh Corgi",
+"Toy Poodle",
+"Miniature Poodle",
+"Standard Poodle",
+"Mexican hairless dog",
+"grey wolf",
+"Alaskan tundra wolf",
+"red wolf",
+"coyote",
+"dingo",
+"dhole",
+"African wild dog",
+"hyena",
+"red fox",
+"kit fox",
+"Arctic fox",
+"grey fox",
+"tabby cat",
+"tiger cat",
+"Persian cat",
+"Siamese cat",
+"Egyptian Mau",
+"cougar",
+"lynx",
+"leopard",
+"snow leopard",
+"jaguar",
+"lion",
+"tiger",
+"cheetah",
+"brown bear",
+"American black bear",
+"polar bear",
+"sloth bear",
+"mongoose",
+"meerkat",
+"tiger beetle",
+"ladybug",
+"ground beetle",
+"longhorn beetle",
+"leaf beetle",
+"dung beetle",
+"rhinoceros beetle",
+"weevil",
+"fly",
+"bee",
+"ant",
+"grasshopper",
+"cricket",
+"stick insect",
+"cockroach",
+"mantis",
+"cicada",
+"leafhopper",
+"lacewing",
+"dragonfly",
+"damselfly",
+"red admiral",
+"ringlet",
+"monarch butterfly",
+"small white",
+"sulphur butterfly",
+"gossamer-winged butterfly",
+"starfish",
+"sea urchin",
+"sea cucumber",
+"cottontail rabbit",
+"hare",
+"Angora rabbit",
+"hamster",
+"porcupine",
+"fox squirrel",
+"marmot",
+"beaver",
+"guinea pig",
+"common sorrel",
+"zebra",
+"pig",
+"wild boar",
+"warthog",
+"hippopotamus",
+"ox",
+"water buffalo",
+"bison",
+"ram",
+"bighorn sheep",
+"Alpine ibex",
+"hartebeest",
+"impala",
+"gazelle",
+"dromedary",
+"llama",
+"weasel",
+"mink",
+"European polecat",
+"black-footed ferret",
+"otter",
+"skunk",
+"badger",
+"armadillo",
+"three-toed sloth",
+"orangutan",
+"gorilla",
+"chimpanzee",
+"gibbon",
+"siamang",
+"guenon",
+"patas monkey",
+"baboon",
+"macaque",
+"langur",
+"black-and-white colobus",
+"proboscis monkey",
+"marmoset",
+"white-headed capuchin",
+"howler monkey",
+"titi",
+"Geoffroy's spider monkey",
+"common squirrel monkey",
+"ring-tailed lemur",
+"indri",
+"Asian elephant",
+"African bush elephant",
+"red panda",
+"giant panda",
+"snoek",
+"eel",
+"coho salmon",
+"rock beauty",
+"clownfish",
+"sturgeon",
+"garfish",
+"lionfish",
+"pufferfish",
+"abacus",
+"abaya",
+"academic gown",
+"accordion",
+"acoustic guitar",
+"aircraft carrier",
+"airliner",
+"airship",
+"altar",
+"ambulance",
+"amphibious vehicle",
+"analog clock",
+"apiary",
+"apron",
+"waste container",
+"assault rifle",
+"backpack",
+"bakery",
+"balance beam",
+"balloon",
+"ballpoint pen",
+"Band-Aid",
+"banjo",
+"baluster",
+"barbell",
+"barber chair",
+"barbershop",
+"barn",
+"barometer",
+"barrel",
+"wheelbarrow",
+"baseball",
+"basketball",
+"bassinet",
+"bassoon",
+"swimming cap",
+"bath towel",
+"bathtub",
+"station wagon",
+"lighthouse",
+"beaker",
+"military cap",
+"beer bottle",
+"beer glass",
+"bell-cot",
+"bib",
+"tandem bicycle",
+"bikini",
+"ring binder",
+"binoculars",
+"birdhouse",
+"boathouse",
+"bobsleigh",
+"bolo tie",
+"poke bonnet",
+"bookcase",
+"bookstore",
+"bottle cap",
+"bow",
+"bow tie",
+"brass",
+"bra",
+"breakwater",
+"breastplate",
+"broom",
+"bucket",
+"buckle",
+"bulletproof vest",
+"high-speed train",
+"butcher shop",
+"taxicab",
+"cauldron",
+"candle",
+"cannon",
+"canoe",
+"can opener",
+"cardigan",
+"car mirror",
+"carousel",
+"tool kit",
+"carton",
+"car wheel",
+"automated teller machine",
+"cassette",
+"cassette player",
+"castle",
+"catamaran",
+"CD player",
+"cello",
+"mobile phone",
+"chain",
+"chain-link fence",
+"chain mail",
+"chainsaw",
+"chest",
+"chiffonier",
+"chime",
+"china cabinet",
+"Christmas stocking",
+"church",
+"movie theater",
+"cleaver",
+"cliff dwelling",
+"cloak",
+"clogs",
+"cocktail shaker",
+"coffee mug",
+"coffeemaker",
+"coil",
+"combination lock",
+"computer keyboard",
+"confectionery store",
+"container ship",
+"convertible",
+"corkscrew",
+"cornet",
+"cowboy boot",
+"cowboy hat",
+"cradle",
+"crane (machine)",
+"crash helmet",
+"crate",
+"infant bed",
+"Crock Pot",
+"croquet ball",
+"crutch",
+"cuirass",
+"dam",
+"desk",
+"desktop computer",
+"rotary dial telephone",
+"diaper",
+"digital clock",
+"digital watch",
+"dining table",
+"dishcloth",
+"dishwasher",
+"disc brake",
+"dock",
+"dog sled",
+"dome",
+"doormat",
+"drilling rig",
+"drum",
+"drumstick",
+"dumbbell",
+"Dutch oven",
+"electric fan",
+"electric guitar",
+"electric locomotive",
+"entertainment center",
+"envelope",
+"espresso machine",
+"face powder",
+"feather boa",
+"filing cabinet",
+"fireboat",
+"fire engine",
+"fire screen sheet",
+"flagpole",
+"flute",
+"folding chair",
+"football helmet",
+"forklift",
+"fountain",
+"fountain pen",
+"four-poster bed",
+"freight car",
+"French horn",
+"frying pan",
+"fur coat",
+"garbage truck",
+"gas mask",
+"gas pump",
+"goblet",
+"go-kart",
+"golf ball",
+"golf cart",
+"gondola",
+"gong",
+"gown",
+"grand piano",
+"greenhouse",
+"grille",
+"grocery store",
+"guillotine",
+"barrette",
+"hair spray",
+"half-track",
+"hammer",
+"hamper",
+"hair dryer",
+"hand-held computer",
+"handkerchief",
+"hard disk drive",
+"harmonica",
+"harp",
+"harvester",
+"hatchet",
+"holster",
+"home theater",
+"honeycomb",
+"hook",
+"hoop skirt",
+"horizontal bar",
+"horse-drawn vehicle",
+"hourglass",
+"iPod",
+"clothes iron",
+"jack-o'-lantern",
+"jeans",
+"jeep",
+"T-shirt",
+"jigsaw puzzle",
+"pulled rickshaw",
+"joystick",
+"kimono",
+"knee pad",
+"knot",
+"lab coat",
+"ladle",
+"lampshade",
+"laptop computer",
+"lawn mower",
+"lens cap",
+"paper knife",
+"library",
+"lifeboat",
+"lighter",
+"limousine",
+"ocean liner",
+"lipstick",
+"slip-on shoe",
+"lotion",
+"speaker",
+"loupe",
+"sawmill",
+"magnetic compass",
+"mail bag",
+"mailbox",
+"tights",
+"tank suit",
+"manhole cover",
+"maraca",
+"marimba",
+"mask",
+"match",
+"maypole",
+"maze",
+"measuring cup",
+"medicine chest",
+"megalith",
+"microphone",
+"microwave oven",
+"military uniform",
+"milk can",
+"minibus",
+"miniskirt",
+"minivan",
+"missile",
+"mitten",
+"mixing bowl",
+"mobile home",
+"Model T",
+"modem",
+"monastery",
+"monitor",
+"moped",
+"mortar",
+"square academic cap",
+"mosque",
+"mosquito net",
+"scooter",
+"mountain bike",
+"tent",
+"computer mouse",
+"mousetrap",
+"moving van",
+"muzzle",
+"nail",
+"neck brace",
+"necklace",
+"nipple",
+"notebook computer",
+"obelisk",
+"oboe",
+"ocarina",
+"odometer",
+"oil filter",
+"organ",
+"oscilloscope",
+"overskirt",
+"bullock cart",
+"oxygen mask",
+"packet",
+"paddle",
+"paddle wheel",
+"padlock",
+"paintbrush",
+"pajamas",
+"palace",
+"pan flute",
+"paper towel",
+"parachute",
+"parallel bars",
+"park bench",
+"parking meter",
+"passenger car",
+"patio",
+"payphone",
+"pedestal",
+"pencil case",
+"pencil sharpener",
+"perfume",
+"Petri dish",
+"photocopier",
+"plectrum",
+"Pickelhaube",
+"picket fence",
+"pickup truck",
+"pier",
+"piggy bank",
+"pill bottle",
+"pillow",
+"ping-pong ball",
+"pinwheel",
+"pirate ship",
+"pitcher",
+"hand plane",
+"planetarium",
+"plastic bag",
+"plate rack",
+"plow",
+"plunger",
+"Polaroid camera",
+"pole",
+"police van",
+"poncho",
+"billiard table",
+"soda bottle",
+"pot",
+"potter's wheel",
+"power drill",
+"prayer rug",
+"printer",
+"prison",
+"projectile",
+"projector",
+"hockey puck",
+"punching bag",
+"purse",
+"quill",
+"quilt",
+"race car",
+"racket",
+"radiator",
+"radio",
+"radio telescope",
+"rain barrel",
+"recreational vehicle",
+"reel",
+"reflex camera",
+"refrigerator",
+"remote control",
+"restaurant",
+"revolver",
+"rifle",
+"rocking chair",
+"rotisserie",
+"eraser",
+"rugby ball",
+"ruler",
+"running shoe",
+"safe",
+"safety pin",
+"salt shaker",
+"sandal",
+"sarong",
+"saxophone",
+"scabbard",
+"weighing scale",
+"school bus",
+"schooner",
+"scoreboard",
+"CRT screen",
+"screw",
+"screwdriver",
+"seat belt",
+"sewing machine",
+"shield",
+"shoe store",
+"shoji",
+"shopping basket",
+"shopping cart",
+"shovel",
+"shower cap",
+"shower curtain",
+"ski",
+"ski mask",
+"sleeping bag",
+"slide rule",
+"sliding door",
+"slot machine",
+"snorkel",
+"snowmobile",
+"snowplow",
+"soap dispenser",
+"soccer ball",
+"sock",
+"solar thermal collector",
+"sombrero",
+"soup bowl",
+"space bar",
+"space heater",
+"space shuttle",
+"spatula",
+"motorboat",
+"spider web",
+"spindle",
+"sports car",
+"spotlight",
+"stage",
+"steam locomotive",
+"through arch bridge",
+"steel drum",
+"stethoscope",
+"scarf",
+"stone wall",
+"stopwatch",
+"stove",
+"strainer",
+"tram",
+"stretcher",
+"couch",
+"stupa",
+"submarine",
+"suit",
+"sundial",
+"sunglass",
+"sunglasses",
+"sunscreen",
+"suspension bridge",
+"mop",
+"sweatshirt",
+"swimsuit",
+"swing",
+"switch",
+"syringe",
+"table lamp",
+"tank",
+"tape player",
+"teapot",
+"teddy bear",
+"television",
+"tennis ball",
+"thatched roof",
+"front curtain",
+"thimble",
+"threshing machine",
+"throne",
+"tile roof",
+"toaster",
+"tobacco shop",
+"toilet seat",
+"torch",
+"totem pole",
+"tow truck",
+"toy store",
+"tractor",
+"semi-trailer truck",
+"tray",
+"trench coat",
+"tricycle",
+"trimaran",
+"tripod",
+"triumphal arch",
+"trolleybus",
+"trombone",
+"tub",
+"turnstile",
+"typewriter keyboard",
+"umbrella",
+"unicycle",
+"upright piano",
+"vacuum cleaner",
+"vase",
+"vault",
+"velvet",
+"vending machine",
+"vestment",
+"viaduct",
+"violin",
+"volleyball",
+"waffle iron",
+"wall clock",
+"wallet",
+"wardrobe",
+"military aircraft",
+"sink",
+"washing machine",
+"water bottle",
+"water jug",
+"water tower",
+"whiskey jug",
+"whistle",
+"wig",
+"window screen",
+"window shade",
+"Windsor tie",
+"wine bottle",
+"wing",
+"wok",
+"wooden spoon",
+"wool",
+"split-rail fence",
+"shipwreck",
+"yawl",
+"yurt",
+"website",
+"comic book",
+"crossword",
+"traffic sign",
+"traffic light",
+"dust jacket",
+"menu",
+"plate",
+"guacamole",
+"consomme",
+"hot pot",
+"trifle",
+"ice cream",
+"ice pop",
+"baguette",
+"bagel",
+"pretzel",
+"cheeseburger",
+"hot dog",
+"mashed potato",
+"cabbage",
+"broccoli",
+"cauliflower",
+"zucchini",
+"spaghetti squash",
+"acorn squash",
+"butternut squash",
+"cucumber",
+"artichoke",
+"bell pepper",
+"cardoon",
+"mushroom",
+"Granny Smith",
+"strawberry",
+"orange",
+"lemon",
+"fig",
+"pineapple",
+"banana",
+"jackfruit",
+"custard apple",
+"pomegranate",
+"hay",
+"carbonara",
+"chocolate syrup",
+"dough",
+"meatloaf",
+"pizza",
+"pot pie",
+"burrito",
+"red wine",
+"espresso",
+"cup",
+"eggnog",
+"alp",
+"bubble",
+"cliff",
+"coral reef",
+"geyser",
+"lakeshore",
+"promontory",
+"shoal",
+"seashore",
+"valley",
+"volcano",
+"baseball player",
+"bridegroom",
+"scuba diver",
+"rapeseed",
+"daisy",
+"yellow lady's slipper",
+"corn",
+"acorn",
+"rose hip",
+"horse chestnut seed",
+"coral fungus",
+"agaric",
+"gyromitra",
+"stinkhorn mushroom",
+"earth star",
+"hen-of-the-woods",
+"bolete",
+"ear of corn",
+"toilet paper"]

pre_run.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+# HF Spaces pre-run hook — installs mmcv-full with the correct CUDA/torch wheel
+set -e
+echo "Installing mmcv-full via openmim..."
+mim install mmcv-full==1.7.2
+echo "mmcv-full installed successfully."

requirements.txt

@@ -0,0 +1,15 @@
+# Core — classification tab (always works)
+torch>=1.12.0
+torchvision>=0.13.0
+timm==0.6.12
+gradio>=4.0.0
+huggingface_hub>=0.20.0
+Pillow>=10.0.0
+numpy>=1.23.0
+scipy
+
+# Detection & Medical tabs
+openmim
+mmdet==2.28.0
+pycocotools
+opencv-python-headless