ecmodels.py

"""
Model wrappers + inference + visualization for EdgeCrafter (ECDet / ECSeg / ECPose).

Ported 1:1 from the authors' reference code:
  - model wiring: EdgeCrafter/hf_models.ipynb (PyTorchModelHubMixin classes)
  - preprocessing / postprocessing / drawing:
        EdgeCrafter/ecdetseg/tools/inference/torch_inf.py
        EdgeCrafter/ecpose/tools/inference/torch_inf.py
"""

from dataclasses import dataclass

import cv2
import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as T
from huggingface_hub import PyTorchModelHubMixin
from PIL import Image

from ecdetseg.engine.edgecrafter.decoder import ECTransformer
from ecdetseg.engine.edgecrafter.ecvit import ViTAdapter
from ecdetseg.engine.edgecrafter.hybrid_encoder import HybridEncoder
from ecdetseg.engine.edgecrafter.postprocessor import PostProcessor
from ecpose.engine.edgecrafter.detrpose_postprocesses import DETRPosePostProcessor
from ecpose.engine.edgecrafter.detrpose_transformer import DETRTransformer
from ecpose.engine.edgecrafter.ecvit import ViTAdapter as PoseViTAdapter
from ecpose.engine.edgecrafter.hybrid_encoder import HybridEncoder as PoseHybridEncoder

# --------------------------------------------------------------------------------------
# COCO metadata (from ecdetseg/engine/data/dataset/coco_dataset.py)
# --------------------------------------------------------------------------------------
COCO_CLASSES = [
    'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
    'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat',
    'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack',
    'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
    'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
    'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
    'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
    'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse',
    'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',
    'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
    'toothbrush',
]

COCO_COLORS = [
    (0, 0, 255), (0, 255, 0), (255, 0, 0), (255, 255, 0), (255, 0, 255),
    (0, 255, 255), (128, 0, 0), (0, 128, 0), (0, 0, 128), (128, 128, 0),
    (128, 0, 128), (0, 128, 128), (255, 128, 0), (255, 0, 128), (0, 255, 128),
    (128, 255, 0), (255, 128, 128), (128, 255, 128), (128, 128, 255), (255, 255, 128),
    (255, 128, 255), (128, 255, 255), (192, 0, 0), (0, 192, 0), (0, 0, 192),
    (192, 192, 0), (192, 0, 192), (0, 192, 192), (255, 192, 0), (255, 0, 192),
    (0, 255, 192), (192, 255, 0), (255, 192, 128), (192, 255, 128), (128, 192, 255),
    (255, 128, 192), (128, 255, 192), (192, 128, 255), (255, 192, 192), (192, 255, 192),
    (192, 192, 255), (255, 255, 192), (255, 192, 255), (192, 255, 255), (64, 0, 0),
    (0, 64, 0), (0, 0, 64), (64, 64, 0), (64, 0, 64), (0, 64, 64), (128, 64, 0),
    (128, 0, 64), (0, 128, 64), (64, 128, 0), (128, 64, 128), (64, 128, 128), (128, 128, 64),
    (192, 64, 0), (192, 0, 64), (0, 192, 64), (64, 192, 0), (192, 64, 192), (64, 192, 192),
    (192, 192, 64), (255, 64, 0), (255, 0, 64), (0, 255, 64), (64, 255, 0),
    (255, 64, 128), (64, 255, 128), (128, 64, 255), (255, 128, 64), (128, 255, 64),
    (64, 128, 255), (192, 64, 128), (192, 128, 64), (64, 192, 128), (128, 192, 64),
    (64, 128, 192), (128, 64, 192), (192, 128, 192), (128, 192, 192), (192, 192, 128),
]

# COCO keypoint skeleton (1-based in the standard definition)
COCO_SKELETON = [
    (16, 14), (14, 12), (17, 15), (15, 13), (12, 13),
    (6, 12), (7, 13), (6, 7), (6, 8), (7, 9),
    (8, 10), (9, 11), (2, 3), (1, 2), (1, 3),
    (2, 4), (3, 5), (4, 6), (5, 7),
]
COCO_SKELETON = [(a - 1, b - 1) for a, b in COCO_SKELETON]

COCO_KEYPOINT_NAMES = [
    "nose", "left_eye", "right_eye", "left_ear", "right_ear",
    "left_shoulder", "right_shoulder", "left_elbow", "right_elbow",
    "left_wrist", "right_wrist", "left_hip", "right_hip",
    "left_knee", "right_knee", "left_ankle", "right_ankle",
]

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
EVAL_SIZE = (640, 640)

_BOUNDARY_KERNEL_CACHE: dict[int, np.ndarray] = {}


# --------------------------------------------------------------------------------------
# Models (mirrors EdgeCrafter/hf_models.ipynb)
# --------------------------------------------------------------------------------------
class ECDet(nn.Module, PyTorchModelHubMixin):
    def __init__(self, config):
        super().__init__()
        config = dict(config)
        config["backbone"] = {**config["backbone"], "skip_load_backbone": True}
        self.config = config

        self.backbone = ViTAdapter(**config["backbone"])
        self.encoder = HybridEncoder(**config["encoder"])
        self.decoder = ECTransformer(**config["decoder"])
        self.postprocessor = PostProcessor(**config["postprocessor"])

    def forward(self, x, orig_target_sizes):
        x = self.backbone(x)
        x = self.encoder(x)
        x = self.decoder(x)
        return self.postprocessor(x, orig_target_sizes)


class ECSeg(nn.Module, PyTorchModelHubMixin):
    def __init__(self, config):
        super().__init__()
        config = dict(config)
        config["backbone"] = {**config["backbone"], "skip_load_backbone": True}
        self.config = config

        self.backbone = ViTAdapter(**config["backbone"])
        self.encoder = HybridEncoder(**config["encoder"])
        self.decoder = ECTransformer(**config["decoder"])
        self.postprocessor = PostProcessor(**config["postprocessor"])

    def forward(self, x, orig_target_sizes):
        x = self.backbone(x)
        x = self.encoder(x)
        # ECSeg feeds the highest-resolution encoder feature to the mask head
        # (see engine/edgecrafter/modeling.py::ECSeg.forward)
        x = self.decoder(x, None, x[0])
        return self.postprocessor(x, orig_target_sizes)


class ECPose(nn.Module, PyTorchModelHubMixin):
    def __init__(self, config):
        super().__init__()
        config = dict(config)
        config["backbone"] = {**config["backbone"], "skip_load_backbone": True}
        self.config = config

        self.backbone = PoseViTAdapter(**config["backbone"])
        self.encoder = PoseHybridEncoder(**config["encoder"])
        self.decoder = DETRTransformer(**config["decoder"])
        self.postprocessor = DETRPosePostProcessor(**config["postprocessor"])

    def forward(self, x, orig_target_sizes):
        x = self.backbone(x)
        x = self.encoder(x)
        x = self.decoder(x, None)
        return self.postprocessor(x, orig_target_sizes)


TASKS = {
    "Object Detection": {"cls": ECDet, "repo": "Intellindust/ECDet_{}"},
    "Instance Segmentation": {"cls": ECSeg, "repo": "Intellindust/ECSeg_{}"},
    "Human Pose Estimation": {"cls": ECPose, "repo": "Intellindust/ECPose_{}"},
}
SIZES = ["S", "M", "L", "X"]


def load_model(task: str, size: str, device: str = "cuda"):
    """Build a deployed (re-parameterized, eval-mode) EdgeCrafter model on `device`."""
    spec = TASKS[task]
    model = spec["cls"].from_pretrained(spec["repo"].format(size))
    model.eval()
    # authors' deploy path: re-parameterize conv/bn blocks, switch postprocessor to
    # tensor (deploy) outputs -- see tools/inference/torch_inf.py::build_model
    for m in model.modules():
        if hasattr(m, "convert_to_deploy"):
            m.convert_to_deploy()
    model.postprocessor.deploy()
    return model.to(device)


_TRANSFORMS = T.Compose([
    T.Resize(EVAL_SIZE),
    T.ToTensor(),
    T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])


def preprocess(image: Image.Image, device: str):
    tensor = _TRANSFORMS(image).unsqueeze(0).to(device)
    orig_sizes = torch.tensor([[image.size[0], image.size[1]]], device=device)
    return tensor, orig_sizes


# --------------------------------------------------------------------------------------
# Results + drawing (ported from the authors' torch_inf.py scripts)
# --------------------------------------------------------------------------------------
@dataclass
class Result:
    label: int
    score: float
    box: np.ndarray = None
    mask: np.ndarray = None
    keypoints: np.ndarray = None


def get_class_color(label: int):
    return COCO_COLORS[label % len(COCO_COLORS)]


def get_draw_params(image_shape):
    height, width = image_shape[:2]
    min_side = max(1, min(height, width))
    base = min_side / 640.0
    font_scale = max(0.6, 0.7 * base)
    text_thickness = max(1, int(round(1.4 * base)))
    box_thickness = max(2, int(round(2.0 * base)))
    boundary_thickness = max(2, int(round(2.0 * base)))
    return font_scale, text_thickness, box_thickness, boundary_thickness


def draw_white_boundary_fast(image: np.ndarray, mask: np.ndarray, thickness: int = 2):
    mask_u8 = mask.astype(np.uint8)
    if not np.any(mask_u8):
        return
    k = max(1, int(thickness))
    kernel = _BOUNDARY_KERNEL_CACHE.get(k)
    if kernel is None:
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * k + 1, 2 * k + 1))
        _BOUNDARY_KERNEL_CACHE[k] = kernel
    edge = cv2.morphologyEx(mask_u8, cv2.MORPH_GRADIENT, kernel)
    image[edge > 0] = (255, 255, 255)


def draw_boxes(image: Image.Image, results: list[Result], alpha: float = 0.5) -> Image.Image:
    im_np = np.array(image, copy=True)
    font_scale, text_thickness, box_thickness, boundary_thickness = get_draw_params(im_np.shape)
    font = cv2.FONT_HERSHEY_SIMPLEX

    if results and alpha > 0:
        overlay = im_np.copy()
        any_mask = False
        merged_mask = np.zeros(im_np.shape[:2], dtype=np.uint8)
        for res in results:
            if res.mask is None:
                continue
            mask_bool = res.mask.astype(bool, copy=False)
            if not np.any(mask_bool):
                continue
            any_mask = True
            overlay[mask_bool] = get_class_color(res.label)
            merged_mask[mask_bool] = 1
        if any_mask:
            im_np = cv2.addWeighted(im_np, 1.0 - alpha, overlay, alpha, 0)
            draw_white_boundary_fast(im_np, merged_mask, thickness=boundary_thickness)

    for res in results:
        color_rgb = get_class_color(res.label)
        x1, y1, x2, y2 = res.box.astype(int)
        x1 = max(0, x1)
        y1 = max(0, y1)
        x2 = min(im_np.shape[1] - 1, x2)
        y2 = min(im_np.shape[0] - 1, y2)

        cv2.rectangle(im_np, (x1, y1), (x2, y2), color_rgb, box_thickness)
        text = f"{class_name(res.label)} {res.score:.2f}"
        (tw, th), baseline = cv2.getTextSize(text, font, font_scale, text_thickness)
        text_x = x1
        text_y = max(th + baseline + 2, y1)
        cv2.rectangle(im_np, (text_x, text_y - th - baseline - 4),
                      (text_x + tw + 4, text_y + 2), color_rgb, -1)
        cv2.putText(im_np, text, (text_x + 2, text_y - baseline - 1), font, font_scale,
                    (255, 255, 255), text_thickness, cv2.LINE_AA)

    return Image.fromarray(im_np.astype(np.uint8))


def draw_pose(image: Image.Image, results: list[Result], draw_skeleton: bool = True) -> Image.Image:
    im_np = np.array(image, copy=True)
    height, width = im_np.shape[:2]
    base = max(1, min(height, width)) / 640.0
    font_scale = max(0.6, 0.7 * base)
    text_thickness = max(1, int(round(1.4 * base)))
    point_radius = max(2, int(round(3.0 * base)))
    line_thickness = max(2, int(round(2.0 * base)))

    for res in results:
        kpts = np.asarray(res.keypoints, dtype=np.float32).reshape(-1, 2).astype(np.int32)
        if kpts.shape[0] == 0:
            continue

        for x, y in kpts:
            cv2.circle(im_np, (int(x), int(y)), point_radius, (0, 255, 0), -1)

        if draw_skeleton:
            for a, b in COCO_SKELETON:
                if a < len(kpts) and b < len(kpts):
                    xa, ya = kpts[a]
                    xb, yb = kpts[b]
                    cv2.line(im_np, (int(xa), int(ya)), (int(xb), int(yb)),
                             (255, 128, 0), line_thickness)

        min_xy = np.maximum(np.min(kpts, axis=0), 0)
        cv2.putText(im_np, f"person {res.score:.2f}",
                    (int(min_xy[0]), int(max(min_xy[1] - 5, 12))),
                    cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255),
                    text_thickness, cv2.LINE_AA)

    return Image.fromarray(im_np.astype(np.uint8))


def class_name(label: int) -> str:
    if 0 <= label < len(COCO_CLASSES):
        return COCO_CLASSES[label]
    return str(label)


# --------------------------------------------------------------------------------------
# Inference (ported from ECInferencer / ECPoseInferencer)
# --------------------------------------------------------------------------------------
@torch.no_grad()
def infer(model, task: str, image: Image.Image, threshold: float, device: str = "cuda"):
    tensor, orig_sizes = preprocess(image, device)

    if task == "Human Pose Estimation":
        scores, labels, keypoints = model(tensor, orig_sizes.to(torch.int64))
        keep = scores[0] > threshold
        scs, lbs, kps = scores[0][keep], labels[0][keep], keypoints[0][keep]
        return [
            Result(label=int(lbs[j].item()), score=float(scs[j].item()),
                   keypoints=kps[j].detach().float().cpu().numpy())
            for j in range(len(scs))
        ]

    outputs = model(tensor, orig_sizes)
    if task == "Instance Segmentation":
        labels, boxes, scores, masks = outputs
    else:
        labels, boxes, scores = outputs
        masks = None

    keep = scores[0] > threshold
    lbls, bxs, scs = labels[0][keep], boxes[0][keep], scores[0][keep]

    results = []
    if masks is not None:
        img_w, img_h = image.size
        # same as the reference (bilinear upsample of the mask logits to the original
        # resolution, threshold at 0), but one kept instance at a time so memory stays
        # bounded on large inputs instead of upsampling all 300 queries at once.
        kept_masks = masks[0][keep].float()
        for j in range(len(lbls)):
            m = torch.nn.functional.interpolate(
                kept_masks[j][None, None], size=(img_h, img_w),
                mode="bilinear", align_corners=False,
            )[0, 0]
            results.append(Result(label=int(lbls[j].item()), score=float(scs[j].item()),
                                  box=bxs[j].float().cpu().numpy(),
                                  mask=(m > 0.0).cpu().numpy()))
    else:
        for j in range(len(lbls)):
            results.append(Result(label=int(lbls[j].item()), score=float(scs[j].item()),
                                  box=bxs[j].float().cpu().numpy()))
    return results


def render(image: Image.Image, task: str, results: list[Result]) -> Image.Image:
    if task == "Human Pose Estimation":
        return draw_pose(image, results)
    return draw_boxes(image, results, alpha=0.5 if task == "Instance Segmentation" else 0.0)


def results_table(task: str, results: list[Result]) -> list[list]:
    rows = []
    if task == "Human Pose Estimation":
        for i, r in enumerate(sorted(results, key=lambda r: -r.score), start=1):
            kpts = np.asarray(r.keypoints, dtype=np.float32).reshape(-1, 2)
            x1, y1 = kpts.min(axis=0)
            x2, y2 = kpts.max(axis=0)
            rows.append([i, "person", round(r.score, 3),
                         f"[{x1:.0f}, {y1:.0f}, {x2:.0f}, {y2:.0f}]", int(kpts.shape[0])])
        return rows

    for i, r in enumerate(sorted(results, key=lambda r: -r.score), start=1):
        x1, y1, x2, y2 = r.box.tolist()
        rows.append([i, class_name(r.label), round(r.score, 3),
                     f"[{x1:.0f}, {y1:.0f}, {x2:.0f}, {y2:.0f}]",
                     int(r.mask.sum()) if r.mask is not None else 0])
    return rows