webui/runner.py

import json
from pathlib import Path
from typing import Dict, Any, List, Tuple

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

from ultralytics import YOLO, SAM

from ImageBind.imagebind import data
from ImageBind.imagebind.models import imagebind_model
from ImageBind.imagebind.models.imagebind_model import ModalityType

import open_clip

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 = []
    idxs = []
    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",
        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)

        # files
        self.id2task_name_path = (self.root / id2task_name_file).resolve()
        self.task2prompt_path = (self.root / task2prompt_file).resolve()
        self.yolo_ckpt_path = (self.root / yolo_ckpt).resolve()

        # load task metadata
        self.id2task_name = json.loads(self.id2task_name_path.read_text())
        self.task2prompt = json.loads(self.task2prompt_path.read_text())

        # caches
        self._vlm_cache = {}
        self._yolo_cache = {}
        self._taskclip_cache = {}

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

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

    def _get_yolo(self, ckpt_path: str):
        ckpt_abs = str((self.root / ckpt_path).resolve()) if ckpt_path.startswith(".") else 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 _get_vlm(self, vlm_model: str):
        if vlm_model in self._vlm_cache:
            return self._vlm_cache[vlm_model]

        if vlm_model == "imagebind":
            m = imagebind_model.imagebind_huge(pretrained=True).to(self.device).eval()
            pack = {"kind": "imagebind", "model": m}
        elif 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}
        elif 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}
        else:
            raise ValueError(f"Unknown vlm_model: {vlm_model}")

        self._vlm_cache[vlm_model] = pack
        return pack

    def _encode_vlm(self, vlm_model: str, prompt_use, seg_list, full_img_pil):
        pack = self._get_vlm(vlm_model)

        with torch.inference_mode():
            if pack["kind"] == "imagebind":
                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 = pack["model"](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 = pack["model"](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 = []
        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, has_cross_attention, ckpt_d_model)
        keys = list(state.keys())
        is_hdc = any(k.startswith("ScoreFunction.HDReason.") for k in keys)
        has_cross = any("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,
    ):
        ckpt_abs = str((self.root / ckpt_path).resolve()) if ckpt_path.startswith(".") else 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

        # IMPORTANT: cache key must include cross_attention + score_function
        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 cross_attention against checkpoint (your training differs by family)
        if bool(cross_attention) != bool(ckpt_has_cross):
            raise RuntimeError(
                f"cross_attention mismatch: runtime={cross_attention} but checkpoint has_cross_attention={ckpt_has_cross}. ckpt={ckpt_abs}"
            )

        # Validate d_model against checkpoint
        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:
            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",
    ) -> 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 truth:
        # - default used cross_attention=True
        # - HDC used cross_attention=False
        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, seg_idxs = _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]} vs image_embedding dim={image_embedding.shape[-1]}"
                )

            d_model = int(image_embedding.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)
                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 = {}
                for p in predict_res:
                    if int(p["category_id"]) == 1:
                        c = p["class"]
                        cluster_scores.setdefault(c, []).append(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 p["category_id"] == 1 and p["class"] != 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},
            }