kiiwee/RB_IBDM_ModelZoo

1import os2 3import random4from dataclasses import dataclass5from typing import Any, List, Dict, Optional, Union, Tuple6import cv27import torch8import requests9import numpy as np10from PIL import Image11import matplotlib.pyplot as plt12from transformers import AutoModelForMaskGeneration, AutoProcessor, pipeline13import gradio as gr14import json15 16 17@dataclass18class BoundingBox:19    xmin: int20    ymin: int21    xmax: int22    ymax: int23 24    @property25    def xyxy(self) -> List[float]:26        return [self.xmin, self.ymin, self.xmax, self.ymax]27@dataclass28class DetectionResult:29    score: float30    label: str31    box: BoundingBox32    mask: Optional[np.ndarray] = None33 34    @classmethod35    def from_dict(cls, detection_dict: Dict) -> 'DetectionResult':36        return cls(37            score=detection_dict['score'],38            label=detection_dict['label'],39            box=BoundingBox(40                xmin=detection_dict['box']['xmin'],41                ymin=detection_dict['box']['ymin'],42                xmax=detection_dict['box']['xmax'],43                ymax=detection_dict['box']['ymax']44            )45        )46 47def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[DetectionResult], include_bboxes: bool = True) -> np.ndarray:48    image_cv2 = np.array(image) if isinstance(image, Image.Image) else image49    image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR)50 51    for detection in detection_results:52        label = detection.label53        score = detection.score54        box = detection.box55        mask = detection.mask56 57        if include_bboxes:58            color = np.random.randint(0, 256, size=3).tolist()59            cv2.rectangle(image_cv2, (box.xmin, box.ymin),60                          (box.xmax, box.ymax), color, 2)61            cv2.putText(image_cv2, f'{label}: {score:.2f}', (box.xmin, box.ymin - 10),62                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)63 64    return cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)65 66 67def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[DetectionResult], include_bboxes: bool = True) -> np.ndarray:68    annotated_image = annotate(image, detections, include_bboxes)69    return annotated_image70 71 72def load_image(image: Union[str, Image.Image]) -> Image.Image:73    if isinstance(image, str) and image.startswith("http"):74        image = Image.open(requests.get(image, stream=True).raw).convert("RGB")75    elif isinstance(image, str):76        image = Image.open(image).convert("RGB")77    else:78        image = image.convert("RGB")79    return image80 81 82def get_boxes(detection_results: List[DetectionResult]) -> List[List[List[float]]]:83    boxes = []84    for result in detection_results:85        xyxy = result.box.xyxy86        boxes.append(xyxy)87    return [boxes]88 89 90def mask_to_polygon(mask: np.ndarray) -> np.ndarray:91    contours, _ = cv2.findContours(92        mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)93    if len(contours) == 0:94        return np.array([])95    largest_contour = max(contours, key=cv2.contourArea)96    return largest_contour97 98 99def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> List[np.ndarray]:100    masks = masks.cpu().float().permute(0, 2, 3, 1).mean(101        axis=-1).numpy().astype(np.uint8)102    masks = (masks > 0).astype(np.uint8)103    if polygon_refinement:104        for idx, mask in enumerate(masks):105            shape = mask.shape106            polygon = mask_to_polygon(mask)107            masks[idx] = cv2.fillPoly(108                np.zeros(shape, dtype=np.uint8), [polygon], 1)109    return list(masks)110 111 112def detect(image: Image.Image, labels: List[str], threshold: float = 0.3, detector_id: Optional[str] = None) -> List[Dict[str, Any]]:113    detector_id = detector_id if detector_id else "IDEA-Research/grounding-dino-base"114    object_detector = pipeline(115        model=detector_id, task="zero-shot-object-detection", device="cpu")116    labels = [label if label.endswith(".") else label+"." for label in labels]117    results = object_detector(118        image, candidate_labels=labels, threshold=threshold)119    return [DetectionResult.from_dict(result) for result in results]120 121 122def segment(image: Image.Image, detection_results: List[DetectionResult], polygon_refinement: bool = False, segmenter_id: Optional[str] = None) -> List[DetectionResult]:123    segmenter_id = segmenter_id if segmenter_id else "martintmv/InsectSAM"124    segmentator = AutoModelForMaskGeneration.from_pretrained(125        segmenter_id).to("cpu")126    processor = AutoProcessor.from_pretrained(segmenter_id)127    boxes = get_boxes(detection_results)128    inputs = processor(images=image, input_boxes=boxes,129                       return_tensors="pt").to("cpu")130    outputs = segmentator(**inputs)131    masks = processor.post_process_masks(132        masks=outputs.pred_masks, original_sizes=inputs.original_sizes, reshaped_input_sizes=inputs.reshaped_input_sizes)[0]133    masks = refine_masks(masks, polygon_refinement)134    for detection_result, mask in zip(detection_results, masks):135        detection_result.mask = mask136    return detection_results137 138 139def grounded_segmentation(image: Union[Image.Image, str], labels: List[str], threshold: float = 0.3, polygon_refinement: bool = False, detector_id: Optional[str] = None, segmenter_id: Optional[str] = None) -> Tuple[np.ndarray, List[DetectionResult]]:140    image = load_image(image)141    detections = detect(image, labels, threshold, detector_id)142    detections = segment(image, detections, polygon_refinement, segmenter_id)143    return np.array(image), detections144 145 146def mask_to_min_max(mask: np.ndarray) -> Tuple[int, int, int, int]:147    y, x = np.where(mask)148    return x.min(), y.min(), x.max(), y.max()149 150 151def extract_and_paste_insect(original_image: np.ndarray, detection: DetectionResult, background: np.ndarray) -> None:152    mask = detection.mask153    xmin, ymin, xmax, ymax = mask_to_min_max(mask)154    insect_crop = original_image[ymin:ymax, xmin:xmax]155    mask_crop = mask[ymin:ymax, xmin:xmax]156 157    insect = cv2.bitwise_and(insect_crop, insect_crop, mask=mask_crop)158 159    x_offset, y_offset = xmin, ymin160    x_end, y_end = x_offset + insect.shape[1], y_offset + insect.shape[0]161 162    insect_area = background[y_offset:y_end, x_offset:x_end]163    insect_area[mask_crop == 1] = insect[mask_crop == 1]164 165 166def create_yellow_background_with_insects(image: np.ndarray) -> np.ndarray:167    labels = ["insect"]168 169    original_image, detections = grounded_segmentation(170        image, labels, threshold=0.3, polygon_refinement=True)171 172    yellow_background = np.full(173        (original_image.shape[0], original_image.shape[1], 3), (0, 255, 255), dtype=np.uint8)  # BGR for yellow174    for detection in detections:175        if detection.mask is not None:176            extract_and_paste_insect(177                original_image, detection, yellow_background)178    # Convert back to RGB to match Gradio's expected input format179    yellow_background = cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)180    return yellow_background181 182 183def run_length_encoding(mask):184    pixels = mask.flatten()185    rle = []186    last_val = 0187    count = 0188    for pixel in pixels:189        if pixel == last_val:190            count += 1191        else:192            if count > 0:193                rle.append(count)194            count = 1195            last_val = pixel196    if count > 0:197        rle.append(count)198    return rle199 200 201def detections_to_json(detections):202    detections_list = []203    for detection in detections:204        detection_dict = {205            "score": detection.score,206            "label": detection.label,207            "box": {208                "xmin": detection.box.xmin,209                "ymin": detection.box.ymin,210                "xmax": detection.box.xmax211            },212            "mask": run_length_encoding(detection.mask) if detection.mask is not None else None213        }214        detections_list.append(detection_dict)215    return detections_list216 217 218def crop_bounding_boxes_with_yellow_background(image: np.ndarray, yellow_background: np.ndarray, detections: List[DetectionResult]) -> List[np.ndarray]:219    crops = []220    for detection in detections:221        xmin, ymin, xmax, ymax = detection.box.xyxy222        crop = yellow_background[ymin:ymax, xmin:xmax]223        crops.append(crop)224    return crops225