Upload 11 files

Dockerfile

@@ -0,0 +1,55 @@
+# Base image: PyTorch with CUDA (torch is pre-installed in conda here)
+FROM pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime
+
+# Environment
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+
+# ---------------------------------------------------------------
+# ALL installs run as ROOT so pip uses the conda environment
+# where torch already lives. This prevents GroundingDINO's
+# setup.py from failing to find torch.
+# ---------------------------------------------------------------
+
+# System dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    ffmpeg \
+    libsm6 \
+    libxext6 \
+    wget \
+    && rm -rf /var/lib/apt/lists/*
+
+# Upgrade pip inside conda
+RUN pip install --no-cache-dir --upgrade pip
+
+# Python dependencies (into the conda env as root)
+COPY requirements.txt /tmp/requirements.txt
+RUN pip install --no-cache-dir -r /tmp/requirements.txt
+
+# Install GroundingDINO with --no-build-isolation so its setup.py can
+# find torch from the conda environment (not a sandboxed temp env).
+RUN pip install --no-cache-dir --no-build-isolation \
+    "git+https://github.com/IDEA-Research/GroundingDINO.git"
+
+# Install Segment Anything
+RUN pip install --no-cache-dir \
+    "git+https://github.com/facebookresearch/segment-anything.git"
+
+# ---------------------------------------------------------------
+# Create non-root user for runtime (HF Spaces requirement)
+# Packages installed above are in /opt/conda and accessible
+# to all users, so this is safe.
+# ---------------------------------------------------------------
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:${PATH}"
+WORKDIR /home/user/app
+
+# Copy application files
+COPY --chown=user . .
+
+# Hugging Face Spaces mandatory port
+EXPOSE 7860
+
+CMD ["python", "app.py"]

app.py

@@ -0,0 +1,92 @@
+"""
+app.py - Hugging Face Spaces entry point.
+Automatically downloads model checkpoints on first run and
+serves a Gradio interface for AI object removal.
+"""
+
+import os
+import sys
+
+# Ensure project root is on the path
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+# -- Step 1: Download checkpoints if needed ----------------------------------
+print("Checking model checkpoints...")
+from download_checkpoints import main as download_checkpoints
+download_checkpoints()
+
+# -- Step 2: Load the pipeline -----------------------------------------------
+import gradio as gr
+from pipeline import ObjectRemovalPipeline
+from config import CLIP_SIMILARITY_THRESHOLD
+
+print("Initializing AI models (GroundingDINO, SAM, CLIP)...")
+pipeline = ObjectRemovalPipeline()
+print("All models loaded. Ready!")
+
+# -- Step 3: Define Gradio callback ------------------------------------------
+def predict(scene_img_path, object_img_paths, threshold):
+    if scene_img_path is None:
+        return None, "Error: Please upload a scene image."
+    if not object_img_paths:
+        return None, "Error: Please upload at least one object reference image."
+
+    # object_img_paths is a list of file paths from gr.File
+    if isinstance(object_img_paths, str):
+        object_img_paths = [object_img_paths]
+
+    try:
+        result_pil = pipeline.run(
+            scene_path=scene_img_path,
+            object_paths=object_img_paths,
+            threshold=threshold,
+            save_debug=False  # avoid cluttering cloud storage
+        )
+        return result_pil, "Done! Object(s) removed successfully."
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        return None, f"Error: {str(e)}"
+
+# -- Step 4: Build Gradio UI -------------------------------------------------
+with gr.Blocks(title="AI Object Eraser") as demo:
+    gr.Markdown("# AI Object Eraser")
+    gr.Markdown(
+        "Upload a **scene photo** and one or more **reference photos** of the objects "
+        "you want to remove. The AI will detect and erase them for you."
+    )
+
+    with gr.Row():
+        with gr.Column():
+            scene_input = gr.Image(
+                label="Scene Image",
+                type="filepath"
+            )
+            object_input = gr.File(
+                label="Object Reference(s) to Remove",
+                file_count="multiple",
+                type="filepath"
+            )
+            threshold_slider = gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=CLIP_SIMILARITY_THRESHOLD,
+                step=0.05,
+                label="Detection Sensitivity (lower = more detections)"
+            )
+            run_btn = gr.Button("Remove Objects", variant="primary")
+
+        with gr.Column():
+            output_img = gr.Image(label="Result")
+            status_box = gr.Textbox(label="Status")
+
+    run_btn.click(
+        fn=predict,
+        inputs=[scene_input, object_input, threshold_slider],
+        outputs=[output_img, status_box]
+    )
+
+# -- Step 5: Launch ----------------------------------------------------------
+if __name__ == "__main__":
+    # Port 7860 is mandatory for Hugging Face Spaces
+    demo.launch(server_name="0.0.0.0", server_port=7860)

clip_matcher.py

@@ -0,0 +1,145 @@
+"""
+models/clip_matcher.py
+-----------------------
+Uses CLIP to compute similarity between a reference-object crop
+and candidate regions in the scene, and to generate text prompts
+for GroundingDINO.
+"""
+
+import os
+from typing import List, Tuple, Optional
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from transformers import CLIPProcessor, CLIPModel
+
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from config import CLIP_MODEL, DEVICE, CLIP_SIMILARITY_THRESHOLD
+
+
+class CLIPMatcher:
+    """
+    Wraps CLIP for two tasks:
+      1. Predict a text description of a reference-object image.
+      2. Score similarity between a reference image and scene crops.
+    """
+
+    CANDIDATE_LABELS = [
+        "a person", "a car", "a chair", "a table", "a dog", "a cat",
+        "a bottle", "a cup", "a book", "a laptop", "a phone", "a bag",
+        "a bicycle", "a motorcycle", "a bus", "a truck", "a tree",
+        "a flower", "a ball", "a clock", "a lamp", "a vase",
+        "a backpack", "furniture", "an electronic device", "a toy",
+        "a plant", "a statue", "a sign", "a box", "an object",
+        # plumbing & household
+        "a tap", "a faucet", "a pipe", "a bucket", "a hose", "a valve",
+        # waste / clutter
+        "trash", "garbage", "waste", "litter", "rubbish",
+        "a plastic bag", "a polythene bag", "a leaf", "dry leaves",
+        "crumpled paper", "a wrapper", "debris",
+    ]
+
+    def __init__(self) -> None:
+        print("  Loading CLIP ...")
+        self.model = CLIPModel.from_pretrained(CLIP_MODEL).to(DEVICE)
+        self.processor = CLIPProcessor.from_pretrained(CLIP_MODEL)
+        self.model.eval()
+
+    @torch.no_grad()
+    def predict_label(self, ref_image: Image.Image) -> str:
+        """
+        Zero-shot classify the reference object and return the best matching
+        text label (usable as a GroundingDINO prompt).
+        """
+        inputs = self.processor(
+            text=self.CANDIDATE_LABELS,
+            images=ref_image,
+            return_tensors="pt",
+            padding=True,
+        ).to(DEVICE)
+
+        outputs = self.model(**inputs)
+        probs = outputs.logits_per_image.softmax(dim=-1)[0]
+        best_idx = probs.argmax().item()
+        best_label = self.CANDIDATE_LABELS[best_idx]
+        best_prob  = probs[best_idx].item()
+
+        print(f"    CLIP label: '{best_label}' (conf={best_prob:.2f})")
+        return best_label
+
+    @torch.no_grad()
+    def image_embedding(self, image: Image.Image) -> torch.Tensor:
+        """Return normalised CLIP image embedding (1 x D)."""
+        inputs = self.processor(images=image, return_tensors="pt").to(DEVICE)
+        emb = self.model.get_image_features(**inputs)
+        if hasattr(emb, "pooler_output"):
+            emb = emb.pooler_output
+        return F.normalize(emb, dim=-1)
+
+    @torch.no_grad()
+    def score_crops(
+        self,
+        ref_image: Image.Image,
+        scene_image: Image.Image,
+        boxes: List[Tuple[int, int, int, int]],
+    ) -> List[float]:
+        """
+        For each box (x1,y1,x2,y2) crop the scene and compute cosine
+        similarity to the reference image embedding.
+        Returns a list of float scores, one per box.
+        """
+        if not boxes:
+            return []
+
+        ref_emb = self.image_embedding(ref_image)          # 1 x D
+
+        crops = []
+        for x1, y1, x2, y2 in boxes:
+            crop = scene_image.crop((x1, y1, x2, y2))
+            crops.append(crop)
+
+        batch_size = 8
+        all_sims = []
+        for i in range(0, len(crops), batch_size):
+            batch_crops = crops[i:i+batch_size]
+            inputs = self.processor(images=batch_crops, return_tensors="pt", padding=True).to(DEVICE)
+            with torch.no_grad():
+                crop_embs = self.model.get_image_features(**inputs)
+            if hasattr(crop_embs, "pooler_output"):
+                crop_embs = crop_embs.pooler_output
+            crop_embs = F.normalize(crop_embs, dim=-1)         # B x D
+            
+            sims = (ref_emb @ crop_embs.T).squeeze(0)          # B
+            
+            # handle 1D vs 0D (if batch_size is 1)
+            if sims.dim() == 0:
+                all_sims.append(sims.item())
+            else:
+                all_sims.extend(sims.cpu().tolist())
+
+        return all_sims
+
+    def filter_boxes_by_similarity(
+        self,
+        ref_image: Image.Image,
+        scene_image: Image.Image,
+        boxes: List[Tuple[int, int, int, int]],
+        threshold: float = CLIP_SIMILARITY_THRESHOLD,
+    ) -> List[Tuple[int, int, int, int]]:
+        """
+        Return only the boxes whose crop has cosine similarity >= threshold
+        with the reference image.
+        """
+        scores = self.score_crops(ref_image, scene_image, boxes)
+        filtered = []
+        for box, score in zip(boxes, scores):
+            print(f"    box {box}  sim={score:.3f}", end="")
+            if score >= threshold:
+                print(" [v]")
+                filtered.append(box)
+            else:
+                print(" [x]")
+        return filtered

config.py

@@ -0,0 +1,75 @@
+"""
+Pipeline configuration - edit these settings to control behaviour.
+"""
+
+import os
+import torch
+
+# -- Device --------------------------------------------------------------------
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE  = torch.float16 if DEVICE == "cuda" else torch.float32
+
+# -- Paths ---------------------------------------------------------------------
+BASE_DIR        = os.path.dirname(os.path.abspath(__file__))
+INPUT_SCENE_DIR = os.path.join(BASE_DIR, "input", "scene")
+INPUT_OBJ_DIR   = os.path.join(BASE_DIR, "input", "objects")
+OUTPUT_DIR      = os.path.join(BASE_DIR, "output")
+CHECKPOINT_DIR  = os.path.join(BASE_DIR, "checkpoints")
+
+# -- GroundingDINO -------------------------------------------------------------
+GDINO_CONFIG = os.path.join(
+    CHECKPOINT_DIR,
+    "GroundingDINO_SwinT_OGC.py"
+)
+GDINO_WEIGHTS = os.path.join(
+    CHECKPOINT_DIR,
+    "groundingdino_swint_ogc.pth"
+)
+GDINO_CONFIG_URL  = (
+    "https://raw.githubusercontent.com/IDEA-Research/GroundingDINO/main/"
+    "groundingdino/config/GroundingDINO_SwinT_OGC.py"
+)
+GDINO_WEIGHTS_URL = (
+    "https://github.com/IDEA-Research/GroundingDINO/releases/download/"
+    "v0.1.0-alpha/groundingdino_swint_ogc.pth"
+)
+
+# -- SAM -----------------------------------------------------------------------
+SAM_CHECKPOINT = os.path.join(CHECKPOINT_DIR, "sam_vit_h_4b8939.pth")
+SAM_MODEL_TYPE = "vit_h"
+SAM_CHECKPOINT_URL = (
+    "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
+)
+
+# -- CLIP (matching) -----------------------------------------------------------
+CLIP_MODEL = "openai/clip-vit-large-patch14"
+
+# -- Inpainting ----------------------------------------------------------------
+# Options: "lama" | "sd_inpaint"
+INPAINT_METHOD = "lama"
+
+# Stable Diffusion inpainting model (used when INPAINT_METHOD == "sd_inpaint")
+SD_INPAINT_MODEL    = "stabilityai/stable-diffusion-2-inpainting"
+SD_INPAINT_STEPS    = 30
+SD_INPAINT_GUIDANCE = 7.5
+SD_INPAINT_PROMPT   = "clean background, sharp, photorealistic, 8k, no blur, seamless"
+
+# LaMa checkpoint (auto-downloaded if missing)
+LAMA_CHECKPOINT = os.path.join(CHECKPOINT_DIR, "big-lama")
+LAMA_URL = (
+    "https://huggingface.co/smartywu/big-lama/resolve/main/big-lama.zip"
+)
+
+# -- Detection / Matching thresholds -------------------------------------------
+# GroundingDINO - box + text confidence
+GDINO_BOX_THRESHOLD  = 0.30
+GDINO_TEXT_THRESHOLD = 0.25
+
+# CLIP cosine similarity threshold for accepting a match
+CLIP_SIMILARITY_THRESHOLD = 0.30
+
+# Mask dilation (pixels) applied before inpainting to cover object edges
+MASK_DILATION_PX = 25
+
+# -- Visualisation -------------------------------------------------------------
+SAVE_DEBUG_IMAGES = True   # saves intermediate masks and detections

detector.py

@@ -0,0 +1,124 @@
+"""
+models/detector.py
+-------------------
+GroundingDINO-based open-vocabulary object detector.
+
+Given a text prompt (generated by CLIP or provided by the user) and a scene
+image, returns bounding boxes with scores.
+"""
+
+import os
+import sys
+from typing import List, Tuple, Dict, Any
+
+import numpy as np
+import torch
+from PIL import Image
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from config import (
+    DEVICE,
+    GDINO_CONFIG, GDINO_WEIGHTS,
+    GDINO_BOX_THRESHOLD, GDINO_TEXT_THRESHOLD,
+)
+
+
+class GroundingDINODetector:
+    """
+    Wraps the GroundingDINO model for open-vocabulary detection.
+    The model is loaded lazily on first use.
+    """
+
+    def __init__(self) -> None:
+        self._model = None
+
+    def _load(self) -> None:
+        if self._model is not None:
+            self._model = self._model.to(DEVICE)
+            return
+
+        print("  Loading GroundingDINO ...")
+        try:
+            from groundingdino.util.inference import load_model
+            self._model = load_model(GDINO_CONFIG, GDINO_WEIGHTS)
+            self._model = self._model.to(DEVICE)
+            self._model.eval()
+        except ImportError as e:
+            raise RuntimeError(
+                "GroundingDINO is not installed. "
+                "Run: pip install groundingdino-py  (or install from source)\n"
+                f"Original error: {e}"
+            )
+
+    @torch.no_grad()
+    def detect(
+        self,
+        image_pil: Image.Image,
+        text_prompt: str,
+        box_threshold:  float = GDINO_BOX_THRESHOLD,
+        text_threshold: float = GDINO_TEXT_THRESHOLD,
+    ) -> List[Dict[str, Any]]:
+        """
+        Run GroundingDINO on `image_pil` with `text_prompt`.
+
+        Returns a list of dicts:
+            { "box": (x1, y1, x2, y2), "score": float, "label": str }
+        All coordinates are in absolute pixels.
+        """
+        self._load()
+
+        from groundingdino.util.inference import predict
+        import groundingdino.datasets.transforms as T
+
+        transform = T.Compose([
+            T.RandomResize([800], max_size=1333),
+            T.ToTensor(),
+            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ])
+
+        image_tensor, _ = transform(image_pil, None)
+        image_tensor = image_tensor.to(DEVICE)
+
+        w, h = image_pil.size
+
+        boxes_norm, scores, labels = predict(
+            model=self._model,
+            image=image_tensor,
+            caption=text_prompt,
+            box_threshold=box_threshold,
+            text_threshold=text_threshold,
+            device=DEVICE,
+        )
+
+        detections = []
+        for box_n, score, label in zip(boxes_norm, scores, labels):
+            # box_n is cx,cy,bw,bh (normalised) -> convert to absolute x1y1x2y2
+            cx, cy, bw, bh = box_n.tolist()
+            x1 = int((cx - bw / 2) * w)
+            y1 = int((cy - bh / 2) * h)
+            x2 = int((cx + bw / 2) * w)
+            y2 = int((cy + bh / 2) * h)
+
+            # clamp to image bounds
+            x1, y1 = max(0, x1), max(0, y1)
+            x2, y2 = min(w, x2), min(h, y2)
+
+            detections.append({
+                "box":   (x1, y1, x2, y2),
+                "score": float(score),
+                "label": label,
+            })
+
+        print(f"    GroundingDINO found {len(detections)} candidate(s) "
+              f"for prompt '{text_prompt}'")
+        return detections
+
+    def boxes_only(
+        self,
+        image_pil: Image.Image,
+        text_prompt: str,
+        **kwargs,
+    ) -> List[Tuple[int, int, int, int]]:
+        """Convenience wrapper - returns only the list of (x1,y1,x2,y2) boxes."""
+        dets = self.detect(image_pil, text_prompt, **kwargs)
+        return [d["box"] for d in dets]

download_checkpoints.py

@@ -0,0 +1,65 @@
+"""
+utils/download_checkpoints.py
+------------------------------
+Downloads all required model checkpoints if not already present.
+Run once before using the pipeline:  python -m utils.download_checkpoints
+"""
+
+import os
+import sys
+
+# allow running standalone
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from config import (
+    CHECKPOINT_DIR,
+    GDINO_CONFIG, GDINO_WEIGHTS, GDINO_CONFIG_URL, GDINO_WEIGHTS_URL,
+    SAM_CHECKPOINT, SAM_CHECKPOINT_URL,
+    LAMA_CHECKPOINT, LAMA_URL,
+    INPAINT_METHOD,
+)
+from image_utils import download_file, download_text_file, unzip
+
+
+def download_groundingdino() -> None:
+    print("\n[1/3] GroundingDINO weights")
+    download_text_file(GDINO_CONFIG_URL, GDINO_CONFIG)
+    download_file(GDINO_WEIGHTS_URL, GDINO_WEIGHTS, "GroundingDINO weights")
+
+
+def download_sam() -> None:
+    print("\n[2/3] SAM (Segment Anything) weights")
+    download_file(SAM_CHECKPOINT_URL, SAM_CHECKPOINT, "SAM ViT-H")
+
+
+def download_lama() -> None:
+    print("\n[3/3] LaMa inpainting weights")
+    if os.path.isdir(LAMA_CHECKPOINT):
+        print("  [DONE] LaMa checkpoint directory already exists")
+        return
+    zip_path = os.path.join(CHECKPOINT_DIR, "big-lama.zip")
+    download_file(LAMA_URL, zip_path, "LaMa big-lama")
+    unzip(zip_path, CHECKPOINT_DIR)
+    os.remove(zip_path)
+    print("  [DONE] LaMa extracted")
+
+
+def main() -> None:
+    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+    print("=" * 60)
+    print("  Downloading model checkpoints")
+    print("=" * 60)
+
+    download_groundingdino()
+    download_sam()
+
+    if INPAINT_METHOD == "lama":
+        download_lama()
+    else:
+        print("\n[3/3] Stable Diffusion -- will auto-download on first run via diffusers")
+
+    print("\n[SUCCESS] All checkpoints ready.\n")
+
+
+if __name__ == "__main__":
+    main()

image_utils.py

@@ -0,0 +1,270 @@
+"""
+utils/image_utils.py
+--------------------
+Image I/O, mask manipulation, and debug-image helpers.
+"""
+
+import os
+import math
+import zipfile
+import requests
+import urllib.request
+from pathlib import Path
+from typing import List, Tuple, Optional, Union
+
+import cv2
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import matplotlib.patches as mpatches
+
+
+# -- Image I/O -----------------------------------------------------------------
+
+def load_image_pil(path: str) -> Image.Image:
+    """Load image as PIL RGB."""
+    return Image.open(path).convert("RGB")
+
+
+def load_image_cv2(path: str) -> np.ndarray:
+    """Load image as OpenCV BGR numpy array."""
+    img = cv2.imread(path)
+    if img is None:
+        raise FileNotFoundError(f"Cannot read image: {path}")
+    return img
+
+
+def pil_to_cv2(img: Image.Image) -> np.ndarray:
+    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+
+
+def cv2_to_pil(img: np.ndarray) -> Image.Image:
+    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+
+
+# Alias for compatibility
+load_image = load_image_pil
+
+
+def show_mask(mask, ax, random_color=False):
+    """Stub for SAM visualization."""
+    pass
+
+
+def show_box(box, ax):
+    """Stub for SAM visualization."""
+    pass
+
+
+def dilate_mask_with_sam_prediction(mask, dilation_px):
+    """Stub for SAM-based dilation."""
+    return mask
+
+
+def save_image(img: Union[Image.Image, np.ndarray], path: str) -> None:
+    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+    if isinstance(img, np.ndarray):
+        cv2.imwrite(path, img)
+    else:
+        img.save(path)
+
+
+def list_images(directory: str) -> List[str]:
+    """Return sorted list of image file paths in a directory."""
+    exts = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"}
+    paths = sorted(
+        str(p) for p in Path(directory).iterdir()
+        if p.suffix.lower() in exts
+    )
+    return paths
+
+
+# -- Mask operations -----------------------------------------------------------
+
+def boxes_to_mask(
+    boxes: List[Tuple[int, int, int, int]],
+    h: int,
+    w: int,
+    dilation_px: int = 0,
+) -> np.ndarray:
+    """
+    Convert list of (x1,y1,x2,y2) boxes to a binary uint8 mask (HW).
+    Optionally dilate the mask by `dilation_px` pixels.
+    """
+    mask = np.zeros((h, w), dtype=np.uint8)
+    for x1, y1, x2, y2 in boxes:
+        x1, y1 = max(0, x1), max(0, y1)
+        x2, y2 = min(w, x2), min(h, y2)
+        mask[y1:y2, x1:x2] = 255
+
+    if dilation_px > 0:
+        kernel = cv2.getStructuringElement(
+            cv2.MORPH_ELLIPSE, (dilation_px * 2 + 1, dilation_px * 2 + 1)
+        )
+        mask = cv2.dilate(mask, kernel)
+
+    return mask
+
+
+def combine_masks(masks: List[np.ndarray]) -> np.ndarray:
+    """OR-combine a list of binary uint8 masks."""
+    if not masks:
+        raise ValueError("Empty mask list")
+    out = np.zeros_like(masks[0])
+    for m in masks:
+        out = cv2.bitwise_or(out, m)
+    return out
+
+
+def refine_mask_with_sam_prediction(
+    raw_mask: np.ndarray,
+    sam_masks: List[np.ndarray],
+) -> np.ndarray:
+    """
+    Given SAM predicted masks (each boolean HW), pick the one with the
+    highest IoU against the raw_mask and return it as uint8.
+    """
+    best_mask = raw_mask
+    best_iou  = 0.0
+    raw_bool  = raw_mask.astype(bool)
+
+    for m in sam_masks:
+        m_bool = m.astype(bool)
+        intersection = (raw_bool & m_bool).sum()
+        union        = (raw_bool | m_bool).sum()
+        iou = intersection / (union + 1e-8)
+        if iou > best_iou:
+            best_iou  = iou
+            best_mask = (m_bool.astype(np.uint8)) * 255
+
+    return best_mask
+
+
+def dilate_mask(mask: np.ndarray, px: int) -> np.ndarray:
+    if px <= 0:
+        return mask
+    kernel = cv2.getStructuringElement(
+        cv2.MORPH_ELLIPSE, (px * 2 + 1, px * 2 + 1)
+    )
+    return cv2.dilate(mask, kernel)
+
+
+# -- Debug visualisation -------------------------------------------------------
+
+def save_detection_debug(
+    scene_path: str,
+    detections: List[dict],
+    output_path: str,
+) -> None:
+    """
+    Draw bounding boxes + labels on the scene image and save.
+    `detections` is a list of dicts with keys: box (x1,y1,x2,y2), label, score.
+    """
+    img = load_image_pil(scene_path)
+    fig, ax = plt.subplots(1, figsize=(12, 8))
+    ax.imshow(img)
+
+    colors = plt.cm.get_cmap("tab10").colors
+    for i, det in enumerate(detections):
+        x1, y1, x2, y2 = det["box"]
+        color = colors[i % len(colors)]
+        rect = mpatches.FancyBboxPatch(
+            (x1, y1), x2 - x1, y2 - y1,
+            boxstyle="round,pad=2",
+            linewidth=2, edgecolor=color, facecolor="none",
+        )
+        ax.add_patch(rect)
+        ax.text(
+            x1, y1 - 6,
+            f"{det['label']} ({det['score']:.2f})",
+            color="white", fontsize=9,
+            bbox=dict(facecolor=color, alpha=0.7, pad=2, edgecolor="none"),
+        )
+
+    ax.axis("off")
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    plt.close()
+
+
+def save_mask_debug(
+    scene_path: str,
+    mask: np.ndarray,
+    output_path: str,
+) -> None:
+    """Overlay the combined mask on the scene image (red, semi-transparent)."""
+    img = np.array(load_image_pil(scene_path))
+    overlay = img.copy()
+    overlay[mask > 0] = [255, 80, 80]
+    blended = cv2.addWeighted(img, 0.55, overlay, 0.45, 0)
+    save_image(Image.fromarray(blended), output_path)
+
+
+def save_comparison(
+    before: Union[Image.Image, np.ndarray],
+    after:  Union[Image.Image, np.ndarray],
+    output_path: str,
+    labels: Tuple[str, str] = ("Before", "After"),
+) -> None:
+    """Save a side-by-side before/after comparison image."""
+    if isinstance(before, np.ndarray):
+        before = cv2_to_pil(before)
+    if isinstance(after, np.ndarray):
+        after = cv2_to_pil(after)
+
+    w  = before.width + after.width + 20
+    h  = max(before.height, after.height) + 40
+    canvas = Image.new("RGB", (w, h), (30, 30, 30))
+    canvas.paste(before, (0, 40))
+    canvas.paste(after,  (before.width + 20, 40))
+
+    # draw labels using matplotlib to avoid font dependency
+    fig, axes = plt.subplots(1, 2, figsize=(14, 7))
+    axes[0].imshow(before); axes[0].set_title(labels[0], fontsize=14); axes[0].axis("off")
+    axes[1].imshow(after);  axes[1].set_title(labels[1], fontsize=14); axes[1].axis("off")
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    plt.close()
+
+
+# -- Checkpoint downloader -----------------------------------------------------
+
+def download_file(url: str, dest: str, desc: str = "") -> None:
+    """Download a file with a progress bar."""
+    os.makedirs(os.path.dirname(dest) or ".", exist_ok=True)
+    if os.path.exists(dest):
+        print(f"  [DONE] Already downloaded: {os.path.basename(dest)}")
+        return
+
+    print(f"  v Downloading {desc or os.path.basename(dest)} ...")
+    response = requests.get(url, stream=True, timeout=120)
+    response.raise_for_status()
+
+    total = int(response.headers.get("content-length", 0))
+    with open(dest, "wb") as f, tqdm(
+        total=total, unit="B", unit_scale=True, desc=desc or os.path.basename(dest)
+    ) as bar:
+        for chunk in response.iter_content(chunk_size=8192):
+            f.write(chunk)
+            bar.update(len(chunk))
+
+
+def download_text_file(url: str, dest: str) -> None:
+    """Download a small text/config file."""
+    os.makedirs(os.path.dirname(dest) or ".", exist_ok=True)
+    if os.path.exists(dest):
+        return
+    print(f"    Fetching config: {os.path.basename(dest)} ...")
+    resp = requests.get(url, timeout=30)
+    resp.raise_for_status()
+    with open(dest, "w") as f:
+        f.write(resp.text)
+
+
+def unzip(zip_path: str, dest_dir: str) -> None:
+    print(f"  -> Extracting {os.path.basename(zip_path)} ...")
+    with zipfile.ZipFile(zip_path, "r") as z:
+        z.extractall(dest_dir)

inpainter_lama.py

@@ -0,0 +1,299 @@
+"""
+inpainter_lama.py  —  big-lama PyTorch inpainter
+Architecture built to match checkpoint indices exactly.
+"""
+
+import os, sys, warnings
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from config import CHECKPOINT_DIR, MASK_DILATION_PX, DEVICE
+
+LAMA_CKPT = os.path.join(CHECKPOINT_DIR, "big-lama", "models", "best.ckpt")
+
+
+# ── FFC primitives ────────────────────────────────────────────────────────────
+
+class FourierUnit(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.conv_layer = nn.Conv2d(in_ch * 2, out_ch * 2, 1, bias=False)
+        self.bn   = nn.BatchNorm2d(out_ch * 2)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        f = torch.fft.rfftn(x, dim=(-2,-1), norm="ortho")
+        f = torch.stack([f.real, f.imag], -1).permute(0,1,4,2,3).contiguous()
+        f = f.view(b, -1, f.shape[-2], f.shape[-1])
+        f = self.relu(self.bn(self.conv_layer(f)))
+        f = f.view(b, -1, 2, f.shape[-2], f.shape[-1]).permute(0,1,3,4,2).contiguous()
+        f = torch.view_as_complex(f)
+        return torch.fft.irfftn(f, s=(h,w), dim=(-2,-1), norm="ortho")
+
+
+class SpectralTransform(nn.Module):
+    def __init__(self, in_ch, out_ch, stride=1):
+        super().__init__()
+        self.downsample = nn.AvgPool2d(2) if stride == 2 else nn.Identity()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_ch, out_ch//2, 1, bias=False),
+            nn.BatchNorm2d(out_ch//2), nn.ReLU(inplace=True))
+        self.fu    = FourierUnit(out_ch//2, out_ch//2)
+        self.conv2 = nn.Conv2d(out_ch//2, out_ch, 1, bias=False)
+
+    def forward(self, x):
+        return self.conv2(self.fu(self.conv1(self.downsample(x))))
+
+
+class FFC(nn.Module):
+    def __init__(self, in_ch, out_ch, ksize, ratio_gin, ratio_gout, stride=1, pad=0):
+        super().__init__()
+        in_cg  = int(in_ch  * ratio_gin);   in_cl  = in_ch  - in_cg
+        out_cg = int(out_ch * ratio_gout);  out_cl = out_ch - out_cg
+        self.in_cg = in_cg
+        self.convl2l = nn.Conv2d(in_cl, out_cl, ksize, stride, pad, bias=False) if in_cl>0 and out_cl>0 else None
+        self.convl2g = nn.Conv2d(in_cl, out_cg, ksize, stride, pad, bias=False) if in_cl>0 and out_cg>0 else None
+        self.convg2l = nn.Conv2d(in_cg, out_cl, ksize, stride, pad, bias=False) if in_cg>0 and out_cl>0 else None
+        self.convg2g = SpectralTransform(in_cg, out_cg, stride)                  if in_cg>0 and out_cg>0 else None
+
+    def forward(self, x):
+        xl, xg = x if isinstance(x, tuple) else (x, None)
+        yl_parts = [f(t) for f, t in [(self.convl2l, xl), (self.convg2l, xg)] if f is not None and t is not None]
+        yg_parts = [f(t) for f, t in [(self.convl2g, xl), (self.convg2g, xg)] if f is not None and t is not None]
+        yl = yl_parts[0] + yl_parts[1] if len(yl_parts) == 2 else (yl_parts[0] if yl_parts else torch.zeros_like(xl))
+        yg = yg_parts[0] + yg_parts[1] if len(yg_parts) == 2 else (yg_parts[0] if yg_parts else None)
+        return yl, yg
+
+
+class FFCBNAct(nn.Module):
+    def __init__(self, in_ch, out_ch, ksize, ratio_gin, ratio_gout, stride=1, pad=0):
+        super().__init__()
+        out_cl = int(out_ch*(1-ratio_gout)); out_cg = out_ch - out_cl
+        self.ffc  = FFC(in_ch, out_ch, ksize, ratio_gin, ratio_gout, stride, pad)
+        self.bn_l = nn.BatchNorm2d(out_cl) if out_cl>0 else None
+        self.bn_g = nn.BatchNorm2d(out_cg) if out_cg>0 else None
+
+    def forward(self, x):
+        xl, xg = self.ffc(x)
+        if self.bn_l: xl = F.relu(self.bn_l(xl), inplace=True)
+        if self.bn_g and xg is not None: xg = F.relu(self.bn_g(xg), inplace=True)
+        return xl, xg
+
+
+class FFCResBlock(nn.Module):
+    def __init__(self, dim, ratio_gin, ratio_gout):
+        super().__init__()
+        self.conv1 = FFCBNAct(dim, dim, 3, ratio_gin,  ratio_gout, pad=1)
+        self.conv2 = FFCBNAct(dim, dim, 3, ratio_gout, ratio_gout, pad=1)
+
+    def forward(self, x):
+        xl, xg = x if isinstance(x, tuple) else (x, None)
+        rl, rg = xl, xg
+        xl, xg = self.conv1((xl, xg))
+        xl, xg = self.conv2((xl, xg))
+        xl = xl + rl
+        if xg is not None and rg is not None: xg = xg + rg
+        return xl, xg
+
+
+# ── Generator — indices match checkpoint exactly ──────────────────────────────
+# Index map from checkpoint:
+#  0 = ReflectionPad2d(3)          [no params]
+#  1 = FFCBNAct(4→64,  k7, 0→0)
+#  2 = FFCBNAct(64→128, k3, 0→0, stride=2)
+#  3 = FFCBNAct(128→256,k3, 0→0, stride=2)
+#  4 = FFCBNAct(256→512,k3, 0→0.75, stride=2)   128 local + 384 global
+#  5-22 = FFCResBlock(512, 0.75→0.75)  × 18
+#  23 = ReLU  [no params]
+#  24 = ConvTranspose2d(512→256)
+#  25 = BatchNorm2d(256)
+#  26 = ReLU  [no params]
+#  27 = ConvTranspose2d(256→128)
+#  28 = BatchNorm2d(128)
+#  29 = ReLU  [no params]
+#  30 = ConvTranspose2d(128→64)
+#  31 = BatchNorm2d(64)
+#  32 = ReLU  [no params]
+#  33 = ReflectionPad2d(3)  [no params]
+#  34 = Conv2d(64→3, k7)
+#  35 = Sigmoid  [no params]
+
+class LaMaGenerator(nn.Module):
+    def __init__(self):
+        super().__init__()
+        ngf = 64
+        self.model = nn.ModuleDict({
+            "1":  FFCBNAct(4,      ngf,   7, 0,    0,    pad=0),
+            "2":  FFCBNAct(ngf,    ngf*2, 3, 0,    0,    stride=2, pad=1),
+            "3":  FFCBNAct(ngf*2,  ngf*4, 3, 0,    0,    stride=2, pad=1),
+            "4":  FFCBNAct(ngf*4,  ngf*8, 3, 0,    0.75, stride=2, pad=1),
+            **{str(i): FFCResBlock(ngf*8, 0.75, 0.75) for i in range(5, 23)},
+            "24": nn.ConvTranspose2d(ngf*8, ngf*4, 3, stride=2, padding=1, output_padding=1),
+            "25": nn.BatchNorm2d(ngf*4),
+            "27": nn.ConvTranspose2d(ngf*4, ngf*2, 3, stride=2, padding=1, output_padding=1),
+            "28": nn.BatchNorm2d(ngf*2),
+            "30": nn.ConvTranspose2d(ngf*2, ngf,   3, stride=2, padding=1, output_padding=1),
+            "31": nn.BatchNorm2d(ngf),
+            "34": nn.Conv2d(ngf, 3, 7, padding=0),
+        })
+
+    def forward(self, x):
+        x = F.pad(x, (3,3,3,3), mode="reflect")          # idx 0
+        x = self.model["1"](x)                            # FFCBNAct → tuple
+        x = self.model["2"](x)
+        x = self.model["3"](x)
+        x = self.model["4"](x)
+        for i in range(5, 23):
+            x = self.model[str(i)](x)
+        # merge local+global
+        xl, xg = x
+        x = torch.cat([xl, xg], dim=1)                   # 512 ch
+        x = F.relu(x, inplace=True)                       # idx 23
+        x = F.relu(self.model["25"](self.model["24"](x)), inplace=True)
+        x = F.relu(self.model["28"](self.model["27"](x)), inplace=True)
+        x = F.relu(self.model["31"](self.model["30"](x)), inplace=True)
+        x = F.pad(x, (3,3,3,3), mode="reflect")          # idx 33
+        x = torch.sigmoid(self.model["34"](x))            # idx 34+35
+        return x
+
+
+# ── Inpainter ─────────────────────────────────────────────────────────────────
+
+class LamaInpainter:
+
+    def __init__(self):
+        self._model   = None
+        self._session = None
+        self._load_pytorch()
+        if self._model is None:
+            self._load_onnx()
+
+    def _load_pytorch(self):
+        if not os.path.exists(LAMA_CKPT):
+            print(f"  big-lama checkpoint not found at {LAMA_CKPT}"); return
+        try:
+            print("  Loading big-lama (PyTorch) ...")
+            ckpt  = torch.load(LAMA_CKPT, map_location="cpu", weights_only=False)
+            state = ckpt.get("state_dict", ckpt)
+            # strip "generator." prefix
+            gen_state = {k[len("generator."):]: v
+                         for k, v in state.items() if k.startswith("generator.")}
+            model = LaMaGenerator()
+            missing, unexpected = model.load_state_dict(gen_state, strict=False)
+            if missing: print(f"  Missing keys: {missing[:3]}")
+            model.eval().to(DEVICE)
+            self._model = model
+            print("  big-lama loaded successfully.")
+        except Exception as e:
+            warnings.warn(f"  big-lama PyTorch load failed: {e}")
+
+    def _load_onnx(self):
+        onnx_path = os.path.join(CHECKPOINT_DIR, "lama_fp32.onnx")
+        try:
+            import onnxruntime as ort
+            from huggingface_hub import hf_hub_download
+            if not os.path.exists(onnx_path):
+                print("  Downloading LaMa ONNX (~125 MB)...")
+                os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+                hf_hub_download(repo_id="Carve/LaMa-ONNX", filename="lama_fp32.onnx",
+                                local_dir=CHECKPOINT_DIR)
+            self._session = ort.InferenceSession(
+                onnx_path, providers=["CUDAExecutionProvider","CPUExecutionProvider"])
+            print("  LaMa ONNX loaded as fallback.")
+        except Exception as e:
+            warnings.warn(f"  LaMa ONNX load failed: {e}. Will use OpenCV.")
+
+    # ── public ────────────────────────────────────────────────────────────────
+
+    def inpaint(self, image_pil: Image.Image, mask: np.ndarray) -> Image.Image:
+        img_np = np.array(image_pil.convert("RGB"))
+        if self._model is not None:
+            import gc
+            try: 
+                return self._pytorch_inpaint(img_np, mask)
+            except Exception as e: 
+                warnings.warn(f"  PyTorch inpaint failed: {e}")
+                import torch
+                torch.cuda.empty_cache()
+                gc.collect()
+
+        # Lazy load ONNX if needed
+        if self._session is None:
+            self._load_onnx()
+
+        if self._session is not None:
+            try: return self._onnx_inpaint(img_np, mask)
+            except Exception as e: warnings.warn(f"  ONNX inpaint failed: {e}")
+
+        return self._opencv_inpaint(img_np, mask)
+
+    def _pytorch_inpaint(self, img_np, mask):
+        h, w = img_np.shape[:2]
+        ys, xs = np.where(mask > 0)
+        if len(xs) == 0: return Image.fromarray(img_np)
+
+        x1,x2 = int(xs.min()), int(xs.max())
+        y1,y2 = int(ys.min()), int(ys.max())
+        pad = max(64, int(max(x2-x1, y2-y1)*0.3))
+        x1=max(0,x1-pad); y1=max(0,y1-pad)
+        x2=min(w,x2+pad); y2=min(h,y2+pad)
+
+        pi = img_np[y1:y2, x1:x2].copy()
+        pm = mask[y1:y2, x1:x2].copy()
+        ph, pw = pi.shape[:2]
+
+        # pad to multiple of 8
+        ph_p = ((ph+7)//8)*8; pw_p = ((pw+7)//8)*8
+        ip = np.pad(pi, ((0,ph_p-ph),(0,pw_p-pw),(0,0)))
+        mp = np.pad(pm, ((0,ph_p-ph),(0,pw_p-pw)))
+
+        it = torch.from_numpy(ip.astype(np.float32)/255.).permute(2,0,1).unsqueeze(0).to(DEVICE)
+        mt = torch.from_numpy((mp>0).astype(np.float32)).unsqueeze(0).unsqueeze(0).to(DEVICE)
+        it = it*(1-mt)
+        inp = torch.cat([it, mt], dim=1)
+
+        with torch.no_grad():
+            out = self._model(inp)
+
+        res = (out[0].permute(1,2,0).cpu().numpy()*255).clip(0,255).astype(np.uint8)[:ph,:pw]
+
+        out_np = img_np.copy()
+        mask_bool = pm > 0
+        out_np[y1:y2, x1:x2][mask_bool] = res[mask_bool]
+        return Image.fromarray(out_np)
+
+    def _onnx_inpaint(self, img_np, mask):
+        SZ = 512
+        h, w = img_np.shape[:2]
+        ys, xs = np.where(mask>0)
+        if len(xs)==0: return Image.fromarray(img_np)
+        x1,x2=int(xs.min()),int(xs.max()); y1,y2=int(ys.min()),int(ys.max())
+        px=max(40,int((x2-x1)*.25)); py=max(40,int((y2-y1)*.25))
+        x1=max(0,x1-px); y1=max(0,y1-py); x2=min(w,x2+px); y2=min(h,y2+py)
+        pi=img_np[y1:y2,x1:x2]; pm=mask[y1:y2,x1:x2]; ph,pw=pi.shape[:2]
+        ir=cv2.resize(pi,(SZ,SZ),interpolation=cv2.INTER_LINEAR)
+        mr=cv2.resize(pm,(SZ,SZ),interpolation=cv2.INTER_NEAREST)
+        it=(ir.astype(np.float32)/255.).transpose(2,0,1)[None]
+        mt=(mr>0).astype(np.float32)[None,None]
+        n0=self._session.get_inputs()[0].name; n1=self._session.get_inputs()[1].name
+        o0=self._session.get_outputs()[0].name
+        res=self._session.run([o0],{n0:it,n1:mt})[0][0]
+        res=(res.transpose(1,2,0)*(255. if res.max()<=1.5 else 1.)).clip(0,255).astype(np.uint8)
+        res=cv2.resize(res,(pw,ph),interpolation=cv2.INTER_LINEAR)
+        out=img_np.copy()
+        mask_bool = pm > 0
+        out[y1:y2,x1:x2][mask_bool] = res[mask_bool]
+        return Image.fromarray(out)
+
+    @staticmethod
+    def _opencv_inpaint(img_np, mask):
+        print("  (OpenCV NS fallback)")
+        bgr=cv2.cvtColor(img_np,cv2.COLOR_RGB2BGR)
+        res=cv2.inpaint(bgr,(mask>0).astype(np.uint8)*255,5,cv2.INPAINT_NS)
+        return Image.fromarray(cv2.cvtColor(res,cv2.COLOR_BGR2RGB))

pipeline.py

@@ -0,0 +1,259 @@
+"""
+pipeline.py
+------------
+Core orchestrator -- ties together detection, segmentation, and inpainting.
+
+Usage (programmatic):
+    from pipeline import ObjectRemovalPipeline
+    pipe = ObjectRemovalPipeline()
+    result = pipe.run(scene_path, object_paths)
+    result.save("output.png")
+
+Usage (CLI):
+    python pipeline.py --scene input/scene/room.jpg \
+                       --objects input/objects/chair.jpg input/objects/lamp.jpg
+"""
+
+import os
+import sys
+import argparse
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import numpy as np
+from PIL import Image
+
+from config import (
+    OUTPUT_DIR, INPAINT_METHOD, SAVE_DEBUG_IMAGES,
+    CLIP_SIMILARITY_THRESHOLD, MASK_DILATION_PX,
+)
+from clip_matcher import CLIPMatcher
+from detector    import GroundingDINODetector
+from segmenter   import SAMSegmenter
+from image_utils  import (
+    load_image_pil, save_image,
+    boxes_to_mask, combine_masks,
+    save_detection_debug, save_mask_debug, save_comparison,
+    list_images,
+)
+
+
+# -- Inpainter factory ---------------------------------------------------------
+
+def _make_inpainter():
+    if INPAINT_METHOD == "lama":
+        from inpainter_lama import LamaInpainter
+        return LamaInpainter()
+    elif INPAINT_METHOD == "sd_inpaint":
+        from inpainter_sd import SDInpainter
+        return SDInpainter()
+    else:
+        raise ValueError(f"Unknown INPAINT_METHOD: {INPAINT_METHOD!r}")
+
+
+# -- Pipeline ------------------------------------------------------------------
+
+class ObjectRemovalPipeline:
+    """
+    End-to-end object removal:
+        reference images  CLIP label  GroundingDINO detection
+         CLIP similarity filter  SAM segmentation  inpainting
+    """
+
+    def __init__(self) -> None:
+        self.clip      = CLIPMatcher()
+        self.detector  = GroundingDINODetector()
+        self.segmenter = SAMSegmenter()
+        self.inpainter = _make_inpainter()
+
+    # -- Public entry point ----------------------------------------------------
+
+    def run(
+        self,
+        scene_path:   str,
+        object_paths: List[str],
+        output_dir:   str = OUTPUT_DIR,
+        save_debug:   bool = SAVE_DEBUG_IMAGES,
+        threshold:    Optional[float] = None,
+    ) -> Image.Image:
+        """
+        Remove all objects specified by `object_paths` from `scene_path`.
+
+        Returns the inpainted PIL image and writes it to `output_dir`.
+        """
+        os.makedirs(output_dir, exist_ok=True)
+        stem = Path(scene_path).stem
+
+        print(f"\n{'='*60}")
+        print(f"  Scene : {scene_path}")
+        print(f"  Objects ({len(object_paths)}): {[os.path.basename(p) for p in object_paths]}")
+        print(f"{'='*60}")
+
+        scene_pil = load_image_pil(scene_path)
+        w, h = scene_pil.size
+
+        # -- Stage 1: detect & match each reference object -----------------
+        all_boxes   = []
+        all_dets    = []   # for debug visualisation
+
+        for obj_path in object_paths:
+            print(f"\n-- Object: {os.path.basename(obj_path)}")
+            obj_pil = load_image_pil(obj_path)
+
+            # 1a. Generate text prompt via CLIP zero-shot classification
+            label = self.clip.predict_label(obj_pil)
+
+            # 1b. Detect in scene with GroundingDINO
+            dets = self.detector.detect(scene_pil, label)
+            candidate_boxes = [d["box"] for d in dets]
+            all_dets.extend(dets)
+
+            if not candidate_boxes:
+                print(f"  [!] No candidates found for '{label}' -- skipping")
+                continue
+
+            # 1c. Filter by CLIP image-image similarity
+            target_threshold = threshold if threshold is not None else CLIP_SIMILARITY_THRESHOLD
+            accepted = self.clip.filter_boxes_by_similarity(
+                obj_pil, scene_pil, candidate_boxes, threshold=target_threshold
+            )
+
+            if not accepted:
+                print(f"  [!] No boxes passed similarity threshold -- skipping")
+            else:
+                print(f"  [v] {len(accepted)} box(es) accepted")
+                all_boxes.extend(accepted)
+
+        if not all_boxes:
+            print("\n[!] No objects detected. Returning original image.")
+            return scene_pil
+
+        # -- Debug: save detection visualisation ---------------------------
+        if save_debug:
+            det_path = os.path.join(output_dir, f"{stem}_debug_detections.jpg")
+            save_detection_debug(scene_path, all_dets, det_path)
+            print(f"\n  [debug] detection image  {det_path}")
+
+        # -- Stage 2: SAM segmentation -------------------------------------
+        print(f"\n-- Segmentation ({len(all_boxes)} box(es))")
+        combined_mask = self.segmenter.segment_boxes(scene_pil, all_boxes)
+
+        if save_debug:
+            mask_path = os.path.join(output_dir, f"{stem}_debug_mask.jpg")
+            save_mask_debug(scene_path, combined_mask, mask_path)
+            print(f"  [debug] mask image      {mask_path}")
+
+        # -- Stage 3: Inpainting -------------------------------------------
+        print(f"\n-- Inpainting ({INPAINT_METHOD})")
+        result_pil = self.inpainter.inpaint(scene_pil, combined_mask)
+
+        # -- Save outputs --------------------------------------------------
+        out_path = os.path.join(output_dir, f"{stem}_result.png")
+        result_pil.save(out_path)
+        print(f"\n  [DONE] Result saved  {out_path}")
+
+        if save_debug:
+            cmp_path = os.path.join(output_dir, f"{stem}_comparison.jpg")
+            save_comparison(scene_pil, result_pil, cmp_path,
+                            labels=("Original", "Objects Removed"))
+            print(f"  [debug] comparison      {cmp_path}")
+
+        return result_pil
+
+    # -- Batch helper ----------------------------------------------------------
+
+    def run_batch(
+        self,
+        scene_dir:   str,
+        objects_dir: str,
+        output_dir:  str = OUTPUT_DIR,
+    ) -> None:
+        """
+        Process all images in `scene_dir`, removing all objects found in
+        `objects_dir`.
+        """
+        scenes  = list_images(scene_dir)
+        objects = list_images(objects_dir)
+
+        if not scenes:
+            print(f"No images found in scene directory: {scene_dir}")
+            return
+        if not objects:
+            print(f"No object images found in: {objects_dir}")
+            return
+
+        print(f"Batch: {len(scenes)} scene(s)  {len(objects)} object reference(s)")
+        for scene_path in scenes:
+            self.run(scene_path, objects, output_dir=output_dir)
+
+
+# -- CLI -----------------------------------------------------------------------
+
+def _parse_args():
+    parser = argparse.ArgumentParser(
+        description="Object Removal Pipeline -- remove specific objects from a scene image."
+    )
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        "--scene", type=str,
+        help="Path to the scene image."
+    )
+    group.add_argument(
+        "--scene-dir", type=str,
+        help="Directory of scene images (batch mode)."
+    )
+    parser.add_argument(
+        "--objects", type=str, nargs="+",
+        help="One or more paths to reference object images."
+    )
+    parser.add_argument(
+        "--objects-dir", type=str,
+        help="Directory of object images (alternative to --objects)."
+    )
+    parser.add_argument(
+        "--output-dir", type=str, default=OUTPUT_DIR,
+        help=f"Output directory (default: {OUTPUT_DIR})."
+    )
+    parser.add_argument(
+        "--inpaint", choices=["lama", "sd_inpaint"],
+        default=None,
+        help="Override inpainting method from config."
+    )
+    parser.add_argument(
+        "--no-debug", action="store_true",
+        help="Skip saving debug images."
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = _parse_args()
+
+    # Override config if flags given
+    if args.inpaint:
+        import config
+        config.INPAINT_METHOD = args.inpaint
+
+    pipe = ObjectRemovalPipeline()
+
+    # Resolve object paths
+    if args.objects:
+        object_paths = args.objects
+    elif args.objects_dir:
+        object_paths = list_images(args.objects_dir)
+    else:
+        print("Error: provide --objects or --objects-dir")
+        sys.exit(1)
+
+    save_debug = not args.no_debug
+
+    if args.scene:
+        pipe.run(args.scene, object_paths,
+                 output_dir=args.output_dir, save_debug=save_debug)
+    elif args.scene_dir:
+        pipe.run_batch(args.scene_dir, args.objects_dir or "",
+                       output_dir=args.output_dir)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,17 @@
+# torch + torchvision are already in the Docker base image (pytorch/pytorch:2.1.2-cuda12.1)
+# Do NOT add torch here - it will conflict with the pre-installed CUDA version
+transformers==4.38.2
+diffusers>=0.24.0
+accelerate>=0.24.0
+opencv-python>=4.8.0
+Pillow>=10.0.0
+numpy>=1.24.0
+scipy>=1.11.0
+scikit-image>=0.21.0
+matplotlib>=3.7.0
+tqdm>=4.66.0
+requests>=2.31.0
+huggingface-hub>=0.19.0
+gradio>=4.0.0
+onnxruntime>=1.16.0
+# segment-anything and groundingdino are installed via git in Dockerfile

segmenter.py

@@ -0,0 +1,133 @@
+"""
+models/segmenter.py
+--------------------
+SAM (Segment Anything Model) wrapper.
+
+Given a scene image and one or more bounding boxes (from the detector),
+produces precise pixel-level masks for each detected object.
+"""
+
+import os
+import sys
+from typing import List, Tuple, Optional
+
+import numpy as np
+import torch
+from PIL import Image
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from config import DEVICE, SAM_CHECKPOINT, SAM_MODEL_TYPE, MASK_DILATION_PX
+from image_utils import load_image, save_image, show_mask, show_box, dilate_mask_with_sam_prediction, dilate_mask, combine_masks
+
+
+class SAMSegmenter:
+    """
+    Wraps SAM to convert bounding boxes into fine-grained masks.
+    Loaded lazily on first use.
+    """
+
+    def __init__(self) -> None:
+        self._predictor = None
+
+    def _load(self) -> None:
+        if self._predictor is not None:
+            self._predictor.model.to(DEVICE)
+            return
+
+        print("  Loading SAM (this may take a moment) ...")
+        try:
+            from segment_anything import sam_model_registry, SamPredictor
+            sam = sam_model_registry[SAM_MODEL_TYPE](checkpoint=SAM_CHECKPOINT)
+            sam = sam.to(DEVICE)
+            self._predictor = SamPredictor(sam)
+        except ImportError as e:
+            raise RuntimeError(
+                "segment-anything is not installed.\n"
+                "Run: pip install git+https://github.com/facebookresearch/segment-anything.git\n"
+                f"Original error: {e}"
+            )
+
+    def segment_boxes(
+        self,
+        image_pil: Image.Image,
+        boxes: List[Tuple[int, int, int, int]],
+        dilation_px: int = MASK_DILATION_PX,
+    ) -> np.ndarray:
+        """
+        For each box, run SAM and return the combined binary mask (HW, uint8).
+
+        Args:
+            image_pil:   The scene image.
+            boxes:       List of (x1, y1, x2, y2) in absolute pixels.
+            dilation_px: How many pixels to dilate the final mask (covers edges).
+
+        Returns:
+            Combined mask (255 = object, 0 = background).
+        """
+        self._load()
+
+        img_np = np.array(image_pil)
+        h, w   = img_np.shape[:2]
+
+        self._predictor.set_image(img_np)
+
+        individual_masks = []
+        for box in boxes:
+            x1, y1, x2, y2 = box
+            box_np = np.array([[x1, y1, x2, y2]], dtype=np.float32)
+
+            masks, scores, _ = self._predictor.predict(
+                box=box_np,
+                multimask_output=True,
+            )
+            # scores shape: (3,); masks shape: (3, H, W)
+            best_idx  = scores.argmax()
+            best_mask = (masks[best_idx].astype(np.uint8)) * 255
+            individual_masks.append(best_mask)
+
+        if not individual_masks:
+            return np.zeros((h, w), dtype=np.uint8)
+
+        combined = combine_masks(individual_masks)
+        if dilation_px > 0:
+            combined = dilate_mask(combined, dilation_px)
+
+        pct = 100 * (combined > 0).sum() / (h * w)
+        print(f"    SAM mask covers {pct:.1f}% of the image")
+        return combined
+
+    def segment_points(
+        self,
+        image_pil: Image.Image,
+        points: List[Tuple[int, int]],
+        point_labels: Optional[List[int]] = None,
+        dilation_px:  int = MASK_DILATION_PX,
+    ) -> np.ndarray:
+        """
+        Segment using foreground point prompts (1 = foreground, 0 = background).
+        Falls back to all-foreground if point_labels is None.
+        """
+        self._load()
+
+        img_np = np.array(image_pil)
+        h, w   = img_np.shape[:2]
+
+        self._predictor.set_image(img_np)
+
+        pts_np    = np.array(points, dtype=np.float32)
+        labels_np = np.array(
+            point_labels if point_labels else [1] * len(points), dtype=np.int32
+        )
+
+        masks, scores, _ = self._predictor.predict(
+            point_coords=pts_np,
+            point_labels=labels_np,
+            multimask_output=True,
+        )
+        best_idx  = scores.argmax()
+        best_mask = (masks[best_idx].astype(np.uint8)) * 255
+
+        if dilation_px > 0:
+            best_mask = dilate_mask(best_mask, dilation_px)
+
+        return best_mask