""" 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