ArtKrit/script/composition/run_models.py

from typing import Any, Dict, List
from PIL import Image
from io import BytesIO
import requests
import numpy as np
from PIL import Image


import torch
import replicate

from .composition_utils import *

try:
    from replicate.helpers import FileOutput as ReplicateFileOutput
except Exception:
    ReplicateFileOutput = None

# ------------------------------
# Model versions
# ------------------------------
REPLICATE_MODEL = "adirik/grounding-dino:efd10a8ddc57ea28773327e881ce95e20cc1d734c589f7dd01d2036921ed78aa"
REPLICATE_SAM_MODEL = "meta/sam-2:fe97b453a6455861e3bac769b441ca1f1086110da7466dbb65cf1eecfd60dc83"

# ------------------------------
# Detector
# ------------------------------
class ReplicateGroundingDetector:
    def __init__(self, model_version: str = REPLICATE_MODEL, box_threshold: float = 0.2, text_threshold: float = 0.2):
        self.model_version = model_version
        self.box_threshold = box_threshold
        self.text_threshold = text_threshold

    def __call__(self, image: Image.Image, candidate_labels: List[str], threshold: float = None) -> List[Dict[str, Any]]:
        import base64

        # Prepare query
        labels = [l if l.endswith(".") else (l + ".") for l in candidate_labels]
        query = " ".join(labels)

        # OPTIMIZATION: Downscale image before uploading to reduce payload size
        max_side = 1280  # Replicate's GroundingDINO works well at this resolution
        W, H = image.size
        if max(W, H) > max_side:
            scale = max_side / max(W, H)
            new_size = (int(W * scale), int(H * scale))
            image_upload = image.resize(new_size, Image.BILINEAR)
            print(f"[GroundingDINO] Downscaling for upload: {(W, H)} -> {new_size}")
        else:
            image_upload = image
            scale = 1.0

        # Convert image to base64 data URI with JPEG (smaller than PNG)
        buffered = BytesIO()
        image_upload.save(buffered, format="JPEG", quality=85)
        img_str = base64.b64encode(buffered.getvalue()).decode()
        data_uri = f"data:image/jpeg;base64,{img_str}"

        # Replicate input
        inputs = {
            "image": data_uri,
            "query": query,
            "box_threshold": threshold if threshold is not None else self.box_threshold,
            "text_threshold": threshold if threshold is not None else self.text_threshold,
        }

        print(f"[Replicate][GroundingDINO] model={self.model_version}\n  query=\"{query}\"\n  box_threshold={inputs['box_threshold']} text_threshold={inputs['text_threshold']}")
        print(f"[Replicate][GroundingDINO] payload size: {len(img_str) / 1024:.1f} KB")

        # Run model with extended timeout
        try:
            out = replicate.run(self.model_version, input=inputs)
        except Exception as e:
            print(f"[Replicate][GroundingDINO] Error: {e}")
            raise

        # Parse detections and scale boxes back to original size
        detections = out.get("detections", [])
        print(f"[Replicate][GroundingDINO] raw detections: {len(detections)}")
        results = []
        for d in detections:
            bbox = d["bbox"]
            item = {
                "score": d.get("score", 0.0),
                "label": d.get("label", ""),
                "box": {
                    "xmin": int(bbox[0] / scale),
                    "ymin": int(bbox[1] / scale),
                    "xmax": int(bbox[2] / scale),
                    "ymax": int(bbox[3] / scale),
                },
            }
            results.append(item)
        if results:
            print("[Replicate][GroundingDINO] parsed detections:")
            for r in results:
                b = r["box"]
                print(f"  - {r['label']} score={r['score']:.2f} box=[{b['xmin']:.1f},{b['ymin']:.1f},{b['xmax']:.1f},{b['ymax']:.1f}]")
        return results

def download_mask(url):
    resp = requests.get(url)
    img = Image.open(BytesIO(resp.content))
    # Prefer alpha channel if present (many mask PNGs encode mask in alpha)
    if "A" in img.getbands():
        mask_pil = img.getchannel("A")
    else:
        mask_pil = img.convert("L")
    mask_np = np.array(mask_pil)
    return mask_np

    
# ------------------------------
# Cloud SAM wrapper
# ------------------------------
def replicate_sam(image_file, boxes=None, **kwargs):
    """
    Cloud-based SAM segmentation using Replicate.
    Accepts a local file-like object (BytesIO or open file).
    If `boxes` is provided, attempt box-prompted segmentation (xyxy in pixel coords).
    """
    # OPTIMIZED: Much lighter configuration to prevent timeouts
    inputs = {
        "image": image_file,
        "use_m2m": False,
        # CRITICAL: Reduce points_per_side dramatically (default is 32!)
        "points_per_side": 4,  # Very aggressive reduction
        # Stricter thresholds = fewer masks
        "pred_iou_thresh": 0.90,
        "stability_score_thresh": 0.92,
        # Additional optimizations
        "crop_n_layers": 0,  # Disable crop-based refinement
        "crop_n_points_downscale_factor": 2,
    }
    # Try common box prompt field names used by Replicate SAM variants
    if boxes is not None:
        inputs["bboxes"] = boxes
        inputs["input_boxes"] = boxes
        inputs["boxes"] = boxes
        inputs["box_format"] = "xyxy"
        inputs["return_individual_masks"] = True
    inputs.update(kwargs)
    print(f"[Replicate][SAM] model={REPLICATE_SAM_MODEL} calling with keys={list(inputs.keys())}")
    
    try:
        out = replicate.run(REPLICATE_SAM_MODEL, input=inputs)
        try:
            if isinstance(out, dict):
                print(f"[Replicate][SAM] returned type={type(out)} keys={list(out.keys())}")
            elif isinstance(out, (list, tuple)):
                print(f"[Replicate][SAM] returned {len(out)} items")
            else:
                print(f"[Replicate][SAM] returned type={type(out)}")
        except Exception as e:
            print(f"[Replicate][SAM] logging error: {e}")
        return out
    except Exception as e:
        print(f"[Replicate][SAM] Error: {e}")
        raise

# ------------------------------
# Initialize models
# ------------------------------
def init_models():
    """
    Initialize the object detector (Replicate Grounding DINO)
    and cloud-based SAM segmenter (Replicate SAM).
    """
    # No local device needed; everything is on Replicate
    object_detector = ReplicateGroundingDetector(
        model_version=REPLICATE_MODEL,
        box_threshold=0.2,
        text_threshold=0.2,
    )

    # Cloud SAM: just a function reference
    segmentator = replicate_sam
    processor = None  # not needed for cloud SAM

    print("✅ Initialized: Using Replicate for detection and segmentation")
    print(f"    - Detector: {REPLICATE_MODEL}")
    print(f"    - SAM:      {REPLICATE_SAM_MODEL}")
    return object_detector, segmentator, processor

# ------------------------------
# Detection pipeline
# ------------------------------
def detect(
    image: Image.Image,
    labels: List[str],
    detector: ReplicateGroundingDetector,
    threshold: float = 0.3,
) -> List[Dict[str, Any]]:
    """
    Detect objects with Replicate Grounding DINO and filter overly large boxes.
    """
    print(f"[Detect] labels={labels} threshold={threshold}")
    raw_results = detector(image, candidate_labels=labels, threshold=threshold)
    image_area = image.size[0] * image.size[1]

    filtered_results = []
    for r in raw_results:
        xmin, ymin, xmax, ymax = r["box"]["xmin"], r["box"]["ymin"], r["box"]["xmax"], r["box"]["ymax"]
        box_area = (xmax - xmin) * (ymax - ymin)
        if image_area > 0 and (box_area / image_area) < 0.8:
            filtered_results.append(DetectionResult.from_dict(r))
    print(f"[Detect] raw={len(raw_results)} filtered={len(filtered_results)}")
    for d in filtered_results:
        print(f"[Detect] keep {d.label} score={d.score:.2f} box={d.box.xyxy}")
    return filtered_results

# ------------------------------
# Segmentation pipeline
# ------------------------------

def _parse_sam_output(out):
    """Normalize various possible Replicate SAM outputs to a list."""
    if isinstance(out, dict):
        print(f"[Segment] SAM dict keys: {list(out.keys())}")
        # Special-case common schema from meta/sam-2 on Replicate
        if "combined_mask" in out:
            ims = out.get("individual_masks")
            parsed = []
            # prefer individual masks first
            if isinstance(ims, list):
                for m in ims:
                    if isinstance(m, dict) and "mask" in m:
                        parsed.append(m["mask"])  # unwrap inner mask field
                    else:
                        parsed.append(m)
            # then include combined if present
            cm = out.get("combined_mask")
            if isinstance(cm, dict) and "mask" in cm:
                parsed.append(cm["mask"])  # unwrap
            elif cm is not None:
                parsed.append(cm)
            return parsed
        for key in ["masks", "mask", "segments", "segmentations", "output", "data"]:
            if key in out:
                return out[key] if isinstance(out[key], list) else [out[key]]
        # fallback: first list-like value
        for v in out.values():
            if isinstance(v, list):
                return v
        return []
    if isinstance(out, (list, tuple)):
        return list(out)
    return [out]


def _coerce_mask_to_numpy(mask_data, target_hw):
    """
    Convert mask outputs (url, data-uri, PIL, ndarray, dict) to a HxW uint8 binary numpy array (0 or 255).
    target_hw = (H, W) of the original image; masks will be resized to this.
    """
    import base64
    H, W = target_hw

    def _ensure_size_u8(m):
        # squeeze channel if needed
        if m.ndim == 3:
            m = m[..., 0]
        if m.shape != (H, W):
            pil = Image.fromarray(m)
            pil = pil.resize((W, H), resample=Image.NEAREST)
            m = np.array(pil)
        # normalize to uint8 binary
        if m.dtype != np.uint8:
            if m.max() <= 1.0:
                m = (m.astype(np.float32) * 255.0).astype(np.uint8)
            else:
                m = m.astype(np.uint8)
        # Many mask PNGs encode binary mask with alpha 0/255; use >0 to be robust
        m = (m > 0).astype(np.uint8) * 255
        return m

    # numpy array
    if isinstance(mask_data, np.ndarray):
        return _ensure_size_u8(mask_data)
    # PIL image
    if isinstance(mask_data, Image.Image):
        return _ensure_size_u8(np.array(mask_data.convert("L")))
    # Replicate FileOutput (URL-like)
    if ReplicateFileOutput is not None and isinstance(mask_data, ReplicateFileOutput):
        try:
            url = getattr(mask_data, "url", None)
            if isinstance(url, str) and url.startswith("http"):
                m = download_mask(url)
                return _ensure_size_u8(m)
        except Exception as e:
            print(f"[Segment] Failed to read Replicate FileOutput: {e}")
            return None
    # string forms
    if isinstance(mask_data, str):
        s = mask_data.strip()
        if s.startswith("http://") or s.startswith("https://"):
            try:
                m = download_mask(s)  # already 0..255 grayscale
                return _ensure_size_u8(m)
            except Exception as e:
                print(f"[Segment] Failed to download mask: {e}")
                return None
        if s.startswith("data:image"):
            try:
                _, b64 = s.split(",", 1)
                img_bytes = base64.b64decode(b64)
                img = Image.open(BytesIO(img_bytes))
                # Prefer alpha channel if present
                if "A" in img.getbands():
                    m = np.array(img.getchannel("A"))
                else:
                    m = np.array(img.convert("L"))
                return _ensure_size_u8(m)
            except Exception as e:
                print(f"[Segment] Failed to decode data URI mask: {e}")
                return None
        # Fallback: some providers return raw base64-encoded PNG without data URI prefix
        try:
            img_bytes = base64.b64decode(s)
            img = Image.open(BytesIO(img_bytes))
            if "A" in img.getbands():
                m = np.array(img.getchannel("A"))
            else:
                m = np.array(img.convert("L"))
            return _ensure_size_u8(m)
        except Exception:
            pass
        print("[Segment] Unknown mask string format; skipping")
        return None
    # dict forms
    if isinstance(mask_data, dict):
        # quick recursive search for any url/data string
        def _find_any_image_string(obj):
            try:
                if isinstance(obj, str) and (obj.startswith("http") or obj.startswith("data:image")):
                    return obj
                if isinstance(obj, dict):
                    for vv in obj.values():
                        s = _find_any_image_string(vv)
                        if s:
                            return s
                if isinstance(obj, (list, tuple)):
                    for vv in obj:
                        s = _find_any_image_string(vv)
                        if s:
                            return s
            except Exception:
                pass
            return None

        # Handle COCO RLE formats (uncompressed only)
        def _decode_coco_rle(rle_obj):
            try:
                if isinstance(rle_obj, dict) and isinstance(rle_obj.get("counts"), list) and "size" in rle_obj:
                    counts = rle_obj["counts"]
                    Hh, Ww = rle_obj["size"]
                    flat = np.zeros(Hh * Ww, dtype=np.uint8)
                    idx = 0
                    val = 0
                    for c in counts:
                        end = idx + int(c)
                        flat[idx:end] = val
                        idx = end
                        val = 255 - val
                    return flat.reshape((Hh, Ww))
                # Compressed RLE (string counts) not supported without pycocotools
                return None
            except Exception as e:
                print(f"[Segment] Failed to decode RLE: {e}")
                return None

        # Top-level RLE
        if "rle" in mask_data and isinstance(mask_data["rle"], (dict,)):
            decoded = _decode_coco_rle(mask_data["rle"])
            if decoded is not None:
                return _ensure_size_u8(decoded)
        # Some schemas put counts/size at top-level
        if "counts" in mask_data and "size" in mask_data:
            decoded = _decode_coco_rle({"counts": mask_data["counts"], "size": mask_data["size"]})
            if decoded is not None:
                return _ensure_size_u8(decoded)

        # Try common fields carrying URLs or data URIs
        for k in ["mask", "url", "image", "overlay", "png", "combined_mask"]:
            v = mask_data.get(k)
            if isinstance(v, str):
                return _coerce_mask_to_numpy(v, target_hw)
            if isinstance(v, dict):
                # nested dict possibly with url or data
                for kk in ["url", "image", "overlay", "data", "png"]:
                    vv = v.get(kk)
                    if isinstance(vv, str):
                        return _coerce_mask_to_numpy(vv, target_hw)
                    # nested numeric array
                    if isinstance(vv, (list, tuple)):
                        arr = np.array(vv)
                        if arr.ndim >= 2:
                            return _ensure_size_u8(arr)
        # If dict has numeric array directly under known keys
        for k in ["data", "array", "segmentation", "mask_array"]:
            v = mask_data.get(k)
            if isinstance(v, (list, tuple)):
                arr = np.array(v)
                if arr.ndim >= 2:
                    return _ensure_size_u8(arr)
        # If dict has a single string value somewhere, try it
        for v in mask_data.values():
            if isinstance(v, str):
                return _coerce_mask_to_numpy(v, target_hw)
        # Final attempt: recursively search any nested url or data-uri string
        s_any = _find_any_image_string(mask_data)
        if s_any:
            return _coerce_mask_to_numpy(s_any, target_hw)
        print("[Segment] Unknown dict mask format; skipping")
        return None
    # list-of-lists (numeric mask)
    if isinstance(mask_data, (list, tuple)):
        try:
            arr = np.array(mask_data)
            if arr.ndim >= 2:
                return _ensure_size_u8(arr)
        except Exception:
            pass
        return None
    # unknown
    return None


def segment(
    image: Image.Image,
    detection_results: List[Any],
    segmentator,
    processor=None,
    device=None,
    polygon_refinement=False,
):
    """
    Segment objects using cloud-based Replicate SAM.
    OPTIMIZED: More aggressive downscaling and fallback strategies.
    """
    W, H = image.size

    # OPTIMIZATION: More aggressive downscaling
    max_side_img = 768  # Reduced from 1024
    scale_img = 1.0
    image_for_sam = image
    if max(W, H) > max_side_img:
        scale_img = max_side_img / float(max(W, H))
        new_size = (int(round(W * scale_img)), int(round(H * scale_img)))
        image_for_sam = image.resize(new_size, Image.BILINEAR)
        print(f"[Segment] Using downscaled image for SAM: {(W, H)} -> {new_size}")

    # Run SAM once with timeout handling
    buf_img = BytesIO()
    image_for_sam.save(buf_img, format="PNG")
    buf_img.seek(0)
    
    try:
        output = segmentator(buf_img)
        parsed = _parse_sam_output(output)
        print(f"[Segment] global SAM outputs={len(parsed)}")
        
        # Coerce all masks to original size (H, W)
        masks_np: List[np.ndarray] = []
        for j, md in enumerate(parsed):
            m = _coerce_mask_to_numpy(md, target_hw=(H, W))
            if m is not None:
                try:
                    nz = int((m > 0).sum())
                    print(f"[Segment] mask[{j}] nz={nz}")
                except Exception:
                    pass
                masks_np.append(m)
    
    except Exception as e:
        print(f"[Segment] SAM failed or timed out: {e}")
        print("[Segment] Falling back to box-fill masks for all detections")
        masks_np = []

    results_with_masks = []
    # Assign best-overlap mask to each detection
    for idx, det in enumerate(detection_results):
        box = det.box
        xmin, ymin, xmax, ymax = map(int, [box.xmin, box.ymin, box.xmax, box.ymax])
        xmin = max(0, min(xmin, W - 1))
        xmax = max(0, min(xmax, W))
        ymin = max(0, min(ymin, H - 1))
        ymax = max(0, min(ymax, H))
        if xmax <= xmin or ymax <= ymin:
            print(f"[Segment] skip invalid box at idx {idx}: {(xmin, ymin, xmax, ymax)}")
            results_with_masks.append(det)
            continue

        box_area = max(1, (xmax - xmin) * (ymax - ymin))
        best_idx = -1
        best_iou = 0.0
        best_metrics = None
        
        # Only try mask matching if we have masks
        if masks_np:
            for k, m in enumerate(masks_np):
                mask_area = int((m > 0).sum())
                # Skip masks that are basically full-frame (likely combined mask)
                if mask_area / float(W * H) > 0.8:
                    continue
                sub = m[ymin:ymax, xmin:xmax]
                overlap = int((sub > 0).sum())
                if overlap == 0:
                    continue
                # IoU with the detection box region
                iou = overlap / float(mask_area + box_area - overlap + 1e-6)
                if iou > best_iou:
                    best_iou = iou
                    best_idx = k
                    best_metrics = (overlap, mask_area)
        
        if best_idx >= 0 and best_iou > 0:
            det.mask = masks_np[best_idx]
            results_with_masks.append(det)
            if idx < 5:
                ov, ma = best_metrics if best_metrics else (0, 0)
                print(f"[Segment] attach mask {best_idx} to det {idx}, overlap={ov}, mask_area={ma}, box_area={box_area}, iou={best_iou:.4f}")
        else:
            print(f"[Segment] no suitable mask for det {idx} — box fill fallback (box_area={box_area})")
            full_mask = np.zeros((H, W), dtype=np.uint8)
            full_mask[ymin:ymax, xmin:xmax] = 255
            det.mask = full_mask
            results_with_masks.append(det)

    return results_with_masks

# ------------------------------
# Device helper
# ------------------------------
def get_device():
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print("CUDA is available. Using GPU.")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
        print("MPS is available! Using Apple GPU.")
    else:
        device = torch.device("cpu")
        print("Using CPU.")
    return device