webui/runner.py

# runner.py
import json
import sys
from pathlib import Path
from typing import Dict, Any, List, Tuple, Optional

from threading import RLock

import numpy as np
import torch
from PIL import Image, ImageDraw

from ultralytics import YOLO, SAM
import open_clip

# --- ImageBind import: robust for both "pip install -e ImageBind" and local folder ---
try:
    # preferred: ImageBind installed as "imagebind"
    from imagebind import data
    from imagebind.models import imagebind_model
    from imagebind.models.imagebind_model import ModalityType
except ModuleNotFoundError:
    # fallback: repo has ./ImageBind/imagebind/
    REPO_ROOT = Path(__file__).resolve().parents[1]  # repo/
    sys.path.insert(0, str(REPO_ROOT / "ImageBind"))
    from imagebind import data
    from imagebind.models import imagebind_model
    from imagebind.models.imagebind_model import ModalityType

from models.TaskCLIP import TaskCLIP


def _draw_boxes_pil(
    img: Image.Image,
    boxes_xyxy: np.ndarray,
    color: Tuple[int, int, int],
    width: int = 3,
) -> Image.Image:
    out = img.copy()
    draw = ImageDraw.Draw(out)
    if boxes_xyxy is None or len(boxes_xyxy) == 0:
        return out
    for (x0, y0, x1, y1) in boxes_xyxy.tolist():
        draw.rectangle([x0, y0, x1, y1], outline=color, width=width)
    return out


def _crop_pil(img: Image.Image, bbox_list: List[List[float]]) -> Tuple[List[Image.Image], List[int]]:
    """Return list of cropped PIL images + indices mapping back to bbox_list."""
    W, H = img.size
    crops: List[Image.Image] = []
    idxs: List[int] = []
    for i, (x0, y0, x1, y1) in enumerate(bbox_list):
        x0 = max(0, min(W, int(x0)))
        y0 = max(0, min(H, int(y0)))
        x1 = max(0, min(W, int(x1)))
        y1 = max(0, min(H, int(y1)))
        if x1 <= x0 or y1 <= y0:
            continue
        crops.append(img.crop((x0, y0, x1, y1)))
        idxs.append(i)
    return crops, idxs


def overlay_masks(
    img: Image.Image,
    masks: np.ndarray,
    alpha: float = 0.40,
    color: Tuple[int, int, int] = (255, 0, 0),
) -> Image.Image:
    if masks is None or len(masks) == 0:
        return img

    base = np.array(img).astype(np.float32)
    union = np.any(masks.astype(bool), axis=0)  # (H, W)
    if not np.any(union):
        return img

    overlay = base.copy()
    overlay[union] = overlay[union] * 0.2 + np.array(color, dtype=np.float32) * 0.8
    out = base * (1 - alpha) + overlay * alpha
    return Image.fromarray(np.clip(out, 0, 255).astype(np.uint8))


class ModelRunner:
    """
    WebUI runner:
      - YOLO detects bboxes
      - VLM (ImageBind or OpenCLIP) embeds text prompts and crops (+ global image)
      - TaskCLIP scores and selects bboxes
      - optionally visualize bbox or SAM masks
    """

    def __init__(
        self,
        project_root: str,
        device: str = "cuda:0",
        yolo_ckpt: str = "./.checkpoints/yolo12x.pt",
        sam_ckpt: str = "./.checkpoints/sam2.1_l.pt",
        imagebind_ckpt: Optional[str] = None,  # optional local weights path
        id2task_name_file: str = "./id2task_name.json",
        task2prompt_file: str = "./task20.json",
        threshold: float = 0.01,
        forward: bool = True,
        cluster: bool = True,
        forward_thre: float = 0.1,
    ):
        self.root = Path(project_root).resolve()
        self.device = device
        self.threshold = float(threshold)
        self.forward = bool(forward)
        self.cluster = bool(cluster)
        self.forward_thre = float(forward_thre)

        # load task metadata
        self.id2task_name_path = (self.root / id2task_name_file).resolve()
        self.task2prompt_path = (self.root / task2prompt_file).resolve()
        self.id2task_name = json.loads(self.id2task_name_path.read_text())
        self.task2prompt = json.loads(self.task2prompt_path.read_text())

        # caches
        self._vlm_cache: Dict[str, Dict[str, Any]] = {}
        self._yolo_cache: Dict[str, YOLO] = {}
        self._taskclip_cache: Dict[Tuple[Any, ...], TaskCLIP] = {}

        # default ckpt paths (not required; YOLO is cached per-run ckpt)
        self.yolo_ckpt_path = (self.root / yolo_ckpt).resolve() if str(yolo_ckpt).startswith(".") else Path(yolo_ckpt)

        # SAM loaded once
        sam_ckpt_path = (self.root / sam_ckpt).resolve() if str(sam_ckpt).startswith(".") else Path(sam_ckpt)
        self.sam = SAM(str(sam_ckpt_path))

        # ImageBind weights path (optional)
        self.imagebind_ckpt = imagebind_ckpt

        # lock for single-GPU servers
        self._lock = RLock()

    def _get_yolo(self, ckpt_path: str) -> YOLO:
        ckpt_abs = str((self.root / ckpt_path).resolve()) if str(ckpt_path).startswith(".") else str(ckpt_path)
        if ckpt_abs not in self._yolo_cache:
            self._yolo_cache[ckpt_abs] = YOLO(ckpt_abs)
        return self._yolo_cache[ckpt_abs]

    def _load_imagebind(self) -> Any:
        """
        Load ImageBind once and cache it.
        - If self.imagebind_ckpt provided and exists: load pretrained=False then load_state_dict
        - Else: pretrained=True (may download)
        """
        if "imagebind" in self._vlm_cache:
            return self._vlm_cache["imagebind"]["model"]

        if self.imagebind_ckpt:
            ckpt_path = (self.root / self.imagebind_ckpt).resolve() if str(self.imagebind_ckpt).startswith(".") else Path(self.imagebind_ckpt)
            if ckpt_path.exists():
                m = imagebind_model.imagebind_huge(pretrained=False).to(self.device).eval()
                state = torch.load(str(ckpt_path), map_location="cpu")
                # common wrappers
                if isinstance(state, dict) and "model" in state and isinstance(state["model"], dict):
                    state = state["model"]
                if isinstance(state, dict) and "state_dict" in state and isinstance(state["state_dict"], dict):
                    state = state["state_dict"]
                m.load_state_dict(state, strict=False)
                self._vlm_cache["imagebind"] = {"kind": "imagebind", "model": m}
                return m

        m = imagebind_model.imagebind_huge(pretrained=True).to(self.device).eval()
        self._vlm_cache["imagebind"] = {"kind": "imagebind", "model": m}
        return m

    def _get_vlm(self, vlm_model: str) -> Dict[str, Any]:
        if vlm_model in self._vlm_cache and vlm_model != "imagebind":
            return self._vlm_cache[vlm_model]

        if vlm_model == "imagebind":
            m = self._load_imagebind()
            return {"kind": "imagebind", "model": m}

        if vlm_model == "vit-b":
            m, _, preprocess = open_clip.create_model_and_transforms(
                "ViT-B-32", pretrained="laion2b_s34b_b79k"
            )
            m = m.to(self.device).eval()
            tokenizer = open_clip.get_tokenizer("ViT-B-32")
            pack = {"kind": "openclip", "model": m, "preprocess": preprocess, "tokenizer": tokenizer}
            self._vlm_cache[vlm_model] = pack
            return pack

        if vlm_model == "vit-l":
            m, _, preprocess = open_clip.create_model_and_transforms(
                "ViT-L-14", pretrained="laion2b_s32b_b82k"
            )
            m = m.to(self.device).eval()
            tokenizer = open_clip.get_tokenizer("ViT-L-14")
            pack = {"kind": "openclip", "model": m, "preprocess": preprocess, "tokenizer": tokenizer}
            self._vlm_cache[vlm_model] = pack
            return pack

        raise ValueError(f"Unknown vlm_model: {vlm_model}")

    def _encode_vlm(
        self,
        vlm_model: str,
        prompt_use: List[str],
        seg_list: List[Image.Image],
        full_img_pil: Image.Image,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        pack = self._get_vlm(vlm_model)

        with torch.inference_mode():
            if pack["kind"] == "imagebind":
                m = pack["model"]
                input_pack = {
                    ModalityType.TEXT: data.load_and_transform_text(prompt_use, self.device),
                    ModalityType.VISION: data.read_and_transform_vision_data(seg_list, self.device),
                }
                emb = m(input_pack)
                text_embeddings = emb[ModalityType.TEXT]
                bbox_embeddings = emb[ModalityType.VISION]

                input_pack2 = {ModalityType.VISION: data.read_and_transform_vision_data([full_img_pil], self.device)}
                emb2 = m(input_pack2)
                image_embedding = emb2[ModalityType.VISION].squeeze(0)
                return text_embeddings, bbox_embeddings, image_embedding

            # openclip branch
            m = pack["model"]
            preprocess = pack["preprocess"]
            tokenizer = pack["tokenizer"]

            # text
            text = tokenizer(prompt_use).to(self.device)
            text_embeddings = m.encode_text(text).float()
            text_embeddings = text_embeddings / text_embeddings.norm(dim=-1, keepdim=True)

            # bbox crops
            crop_tensors = [preprocess(im) for im in seg_list]
            crop_batch = torch.stack(crop_tensors, dim=0).to(self.device)
            bbox_embeddings = m.encode_image(crop_batch).float()
            bbox_embeddings = bbox_embeddings / bbox_embeddings.norm(dim=-1, keepdim=True)

            # global image
            img_tensor = preprocess(full_img_pil).unsqueeze(0).to(self.device)
            image_embedding = m.encode_image(img_tensor).float().squeeze(0)
            image_embedding = image_embedding / image_embedding.norm(dim=-1, keepdim=True)

            return text_embeddings, bbox_embeddings, image_embedding

    def list_task_ids(self) -> List[int]:
        ids: List[int] = []
        for k in self.id2task_name.keys():
            try:
                ids.append(int(k))
            except Exception:
                pass
        return sorted(ids)

    @staticmethod
    def _unwrap_state_dict(obj: Any) -> Dict[str, torch.Tensor]:
        # supports {"state_dict": ...} style checkpoints
        if isinstance(obj, dict) and "state_dict" in obj and isinstance(obj["state_dict"], dict):
            return obj["state_dict"]
        if isinstance(obj, dict):
            return obj
        raise TypeError(f"Unsupported checkpoint format: {type(obj)}")

    def _infer_ckpt_flags(self, state: Dict[str, torch.Tensor]) -> Tuple[bool, bool, int]:
        """
        Infer:
          - is_hdc: whether checkpoint contains HDC submodule keys
          - has_cross_attention: whether checkpoint contains cross-attn keys
          - ckpt_d_model: best-effort inferred d_model from weights
        """
        keys = list(state.keys())
        is_hdc = any(k.startswith("ScoreFunction.HDReason.") for k in keys)
        # NOTE: adjust this if your TaskCLIP names cross-attn differently
        has_cross = any("cross_attn" in k or "cross_attn_text" in k for k in keys)

        if "decoder_norm.weight" in state:
            ckpt_d_model = int(state["decoder_norm.weight"].shape[0])
        elif "ScoreFunction.norm.weight" in state:
            ckpt_d_model = int(state["ScoreFunction.norm.weight"].shape[0])
        else:
            ckpt_d_model = -1

        return is_hdc, has_cross, ckpt_d_model

    def _get_taskclip(
        self,
        ckpt_path: str,
        d_model: int,
        n_words: int,
        score_function: str,
        hdv_dim: int,
        cross_attention: bool,
    ) -> TaskCLIP:
        ckpt_abs = str((self.root / ckpt_path).resolve()) if str(ckpt_path).startswith(".") else str(ckpt_path)
        if not Path(ckpt_abs).exists():
            raise FileNotFoundError(f"TaskCLIP checkpoint not found: {ckpt_abs}")

        eff_hdv_dim = int(hdv_dim) if score_function == "HDC" else 0

        # cache key must include cross_attention + score_function + dimensions
        key = (ckpt_abs, int(d_model), int(n_words), str(score_function), int(eff_hdv_dim), bool(cross_attention))
        if key in self._taskclip_cache:
            return self._taskclip_cache[key]

        state_raw = torch.load(ckpt_abs, map_location="cpu")
        state = self._unwrap_state_dict(state_raw)

        ckpt_is_hdc, ckpt_has_cross, ckpt_d_model = self._infer_ckpt_flags(state)

        # Validate score_function against checkpoint
        if score_function == "HDC" and not ckpt_is_hdc:
            raise RuntimeError(f"Checkpoint is NOT HDC but score_function=HDC was selected. ckpt={ckpt_abs}")
        if score_function != "HDC" and ckpt_is_hdc:
            raise RuntimeError(f"Checkpoint IS HDC but score_function=default was selected. ckpt={ckpt_abs}")

        # Validate d_model against checkpoint (if inferred)
        if ckpt_d_model != -1 and int(d_model) != int(ckpt_d_model):
            raise RuntimeError(
                f"d_model mismatch: VLM produced d_model={int(d_model)} but checkpoint expects d_model={int(ckpt_d_model)}. ckpt={ckpt_abs}"
            )

        model_config = {
            "num_layers": 8,
            "norm": None,
            "return_intermediate": False,
            "d_model": int(d_model),
            "nhead": 4,
            "dim_feedforward": 2048,
            "dropout": 0.1,
            "N_words": int(n_words),
            "activation": "gelu",
            "normalize_before": False,
            "device": self.device,
            "ratio_text": 0.3,
            "ratio_image": 0.3,
            "ratio_glob": 0.3,
            "norm_before": True,
            "norm_after": False,
            "MIN_VAL": 10.0,
            "MAX_VAL": 30.0,
            "cross_attention": bool(cross_attention),
            "score_function": "HDC" if score_function == "HDC" else "default",
            "HDV_D": int(eff_hdv_dim),
        }

        m = TaskCLIP(model_config, normalize_before=model_config["normalize_before"], device=model_config["device"])
        m.load_state_dict(state, strict=True)
        m = m.to(self.device).eval()

        self._taskclip_cache[key] = m
        return m

    def _find_same_class(self, predict_res, score, visited, i, classes, confs, forward_thre):
        cls_i = classes[i]
        for j in range(len(score)):
            if visited[j] == 1:
                continue
            if classes[j] == cls_i and float(score[j]) > forward_thre:
                visited[j] = 1
                predict_res[j]["category_id"] = 1
                predict_res[j]["score"] = float(score[j])

    def _sam_masks_from_bboxes(self, image_path: str, bbox_list: List[List[float]], img_h: int, img_w: int) -> np.ndarray:
        if not bbox_list:
            return np.zeros((0, img_h, img_w), dtype=bool)

        bboxes = [[float(x0), float(y0), float(x1), float(y1)] for x0, y0, x1, y1 in bbox_list]

        try:
            res = self.sam(image_path, bboxes=bboxes)
            r0 = res[0]
            if r0.masks is None:
                return np.zeros((0, img_h, img_w), dtype=bool)
            masks = r0.masks.data.detach().cpu().numpy().astype(bool)
            return masks
        except Exception:
            # fallback per-box
            masks_list = []
            for bb in bboxes:
                rr = self.sam(image_path, bboxes=bb)[0]
                if rr.masks is None:
                    continue
                m = rr.masks.data.detach().cpu().numpy().astype(bool)
                masks_list.append(m[0])
            if len(masks_list) == 0:
                return np.zeros((0, img_h, img_w), dtype=bool)
            return np.stack(masks_list, axis=0)

    def run(
        self,
        image_path: str,
        task_id: int,
        vlm_model: str = "imagebind",
        od_model: str = "yolo",
        yolo_ckpt: str = "./.checkpoints/yolo12x.pt",
        score_function: str = "default",
        hdv_dim: int = 256,
        taskclip_ckpt: str = "./test_model/default/decoder.pt",
        viz_mode: str = "bbox",
        hw_noise_dist: str = "none",
        hw_noise_width: int = 0,
        hw_noise_strength: int = 0,
        hdc_bits: int = 32
    ) -> Dict[str, Any]:
        if vlm_model not in ["imagebind", "vit-b", "vit-l"]:
            raise ValueError(f"Unknown vlm_model: {vlm_model}")
        if od_model != "yolo":
            raise ValueError("Currently only od_model='yolo' is supported.")
        if viz_mode not in ["bbox", "mask"]:
            raise ValueError(f"Unknown viz_mode={viz_mode}")

        # Training convention you stated:
        # - default => cross_attention True
        # - HDC     => cross_attention False
        # If your actual training differs, change this rule OR pass it from app.py.
        cross_attention = (score_function != "HDC")

        with self._lock:
            img = Image.open(image_path).convert("RGB")

            task_name = self.id2task_name[str(task_id)]
            prompt_words = self.task2prompt[task_name]
            prompt_use = ["The item is " + w for w in prompt_words]

            # YOLO detect
            yolo = self._get_yolo(yolo_ckpt)
            outputs = yolo(image_path)
            bbox_list = outputs[0].boxes.xyxy.tolist()
            classes = outputs[0].boxes.cls.tolist()
            confidences = outputs[0].boxes.conf.tolist()

            H, W = img.size[1], img.size[0]
            all_boxes = np.asarray(bbox_list, dtype=np.float32)

            # visualize all detections
            if viz_mode == "bbox":
                img_yolo = _draw_boxes_pil(img, all_boxes, color=(0, 255, 0), width=3)
                all_masks = None
            else:
                all_masks = self._sam_masks_from_bboxes(image_path, bbox_list, img_h=H, img_w=W)
                img_yolo = overlay_masks(img, all_masks, alpha=0.35, color=(0, 255, 0))

            # crop bboxes
            seg_list, _ = _crop_pil(img, bbox_list)
            if len(seg_list) == 0:
                return {
                    "task_id": task_id,
                    "task_name": task_name,
                    "bbox_list": bbox_list,
                    "selected_indices": [],
                    "images": {"original": img, "yolo": img_yolo, "selected": img.copy()},
                }

            # VLM embeddings
            text_embeddings, bbox_embeddings, image_embedding = self._encode_vlm(
                vlm_model=vlm_model,
                prompt_use=prompt_use,
                seg_list=seg_list,
                full_img_pil=img,
            )

            # Ensure dims are consistent
            if int(bbox_embeddings.shape[-1]) != int(image_embedding.shape[-1]):
                raise RuntimeError(
                    f"Embedding dim mismatch: bbox_embeddings dim={bbox_embeddings.shape[-1]} "
                    f"vs image_embedding dim={image_embedding.shape[-1]}"
                )

            # IMPORTANT: d_model should come from bbox_embeddings (tgt), not global image
            d_model = int(bbox_embeddings.shape[-1])
            n_words = int(text_embeddings.shape[0])

            # TaskCLIP (load correct arch)
            taskclip = self._get_taskclip(
                ckpt_path=taskclip_ckpt,
                d_model=d_model,
                n_words=n_words,
                score_function=score_function,
                hdv_dim=hdv_dim,
                cross_attention=cross_attention,
            )

            # Score
            with torch.inference_mode():
                tgt = bbox_embeddings
                memory = text_embeddings
                image_embedding_2d = image_embedding.view(1, -1)
                _, _, score_res, _ = taskclip(
                    tgt, 
                    memory, 
                    image_embedding_2d,
                    hw_noise_dist=hw_noise_dist,
                    hw_noise_width=int(hw_noise_width),
                    hw_noise_strength=int(hw_noise_strength),
                    hdc_bits=hdc_bits)
                score = score_res.view(-1).detach().cpu().numpy().tolist()

            # post-process
            predict_res = []
            for i in range(len(bbox_list)):
                predict_res.append({"category_id": -1, "score": -1, "class": int(classes[i])})

            visited = [0] * len(score)
            for i, x in enumerate(score):
                if visited[i] == 1:
                    continue
                if float(x) > self.threshold:
                    visited[i] = 1
                    predict_res[i]["category_id"] = 1
                    predict_res[i]["score"] = float(x)
                    if self.forward:
                        self._find_same_class(predict_res, score, visited, i, classes, confidences, self.forward_thre)
                else:
                    predict_res[i]["category_id"] = 0
                    predict_res[i]["score"] = 1.0 - float(x)

            # cluster optimization
            if self.cluster and self.forward and len(seg_list) > 1:
                cluster_scores: Dict[int, List[float]] = {}
                for p in predict_res:
                    if int(p["category_id"]) == 1:
                        c = int(p["class"])
                        cluster_scores.setdefault(c, []).append(float(p["score"]))

                if len(cluster_scores) > 1:
                    cluster_ave = {c: float(np.mean(v)) for c, v in cluster_scores.items()}
                    select_class = max(cluster_ave, key=lambda k: cluster_ave[k])
                    for p in predict_res:
                        if int(p["category_id"]) == 1 and int(p["class"]) != int(select_class):
                            p["category_id"] = 0

            selected_indices = [i for i, p in enumerate(predict_res) if int(p["category_id"]) == 1]
            selected_boxes = (
                all_boxes[selected_indices] if len(selected_indices) > 0 else np.zeros((0, 4), dtype=np.float32)
            )

            # visualize selected
            if viz_mode == "bbox":
                img_selected = _draw_boxes_pil(img, selected_boxes, color=(255, 0, 0), width=4)
            else:
                if all_masks is not None and all_masks.shape[0] > 0 and len(selected_indices) > 0:
                    sel_masks = all_masks[selected_indices]
                else:
                    sel_masks = np.zeros((0, H, W), dtype=bool)
                img_selected = overlay_masks(img, sel_masks, alpha=0.45, color=(255, 0, 0))

            return {
                "task_id": task_id,
                "task_name": task_name,
                "bbox_list": bbox_list,
                "classes": classes,
                "confidences": confidences,
                "scores": score,
                "selected_indices": selected_indices,
                "images": {"original": img, "yolo": img_yolo, "selected": img_selected},
            }