Augmentus-robotics/object-analysis-api

1# utils/segment_model.py2from __future__ import annotations3 4import io5import os6import json7import base648from typing import List, Tuple, Union, Optional9import numpy as np10import cv211from PIL import Image as PILImage12from google import genai13from google.genai import types14 15 16 17def parse_json(text: str) -> str:18    return text.strip().removeprefix("```json").removesuffix("```")19 20 21def generate_mask(predicted_str: str, *, img_height: int, img_width: int) -> list[tuple[np.ndarray, str]]:22    try:23        items = json.loads(parse_json(predicted_str))24        if not isinstance(items, list):25            print("Error: Parsed JSON is not a list.")26            return []27    except json.JSONDecodeError as e:28        print(f"Error decoding JSON: {e}")29        print(f"Problematic string snippet: {predicted_str[:200]}...")30        return []31    except Exception as e:32        print(f"An unexpected error occurred during JSON parsing: {e}")33        return []34 35    segmentation_data = []36    default_label = "unknown"37 38    for item_idx, item in enumerate(items):39        if not isinstance(item, dict) or "box_2d" not in item or "mask" not in item:40            print(f"Skipping invalid item structure at index {item_idx}: {item}")41            continue42 43        label = item.get("label", default_label)44        if not isinstance(label, str) or not label:45            label = default_label46 47        png_str = item["mask"]48        if not isinstance(png_str, str) or not png_str.startswith("data:image/png;base64,"):49            print(f"Skipping item {item_idx} (label: {label}) with invalid mask format.")50            continue51        png_str = png_str.removeprefix("data:image/png;base64,")52        try:53            png_bytes = base64.b64decode(png_str)54            bbox_mask = cv2.imdecode(np.frombuffer(png_bytes, np.uint8), cv2.IMREAD_GRAYSCALE)55            if bbox_mask is None:56                print(f"Skipping item {item_idx} (label: {label}) because mask decoding failed.")57                continue58        except (base64.binascii.Error, ValueError, Exception) as e:59            print(f"Error decoding base64 or image data for item {item_idx} (label: {label}): {e}")60            continue61 62        try:63            box = item["box_2d"]64            if not isinstance(box, list) or len(box) != 4:65                print(f"Skipping item {item_idx} (label: {label}) with invalid box_2d format: {box}")66                continue67            y0_norm, x0_norm, y1_norm, x1_norm = map(float, box)68            abs_y0 = max(0, min(int(y0_norm / 1000.0 * img_height), img_height - 1))69            abs_x0 = max(0, min(int(x0_norm / 1000.0 * img_width), img_width - 1))70            abs_y1 = max(0, min(int(y1_norm / 1000.0 * img_height), img_height))71            abs_x1 = max(0, min(int(x1_norm / 1000.0 * img_width), img_width))72            bbox_height = abs_y1 - abs_y073            bbox_width = abs_x1 - abs_x074            if bbox_height <= 0 or bbox_width <= 0:75                print(f"Skipping item {item_idx} (label: {label}) with invalid bbox dims: {box} -> ({bbox_width}x{bbox_height})")76                continue77        except (ValueError, TypeError) as e:78            print(f"Skipping item {item_idx} (label: {label}) due to error processing box_2d: {e}")79            continue80 81        try:82            if bbox_mask.shape[0] > 0 and bbox_mask.shape[1] > 0:83                resized_bbox_mask = cv2.resize(84                    bbox_mask, (bbox_width, bbox_height), interpolation=cv2.INTER_LINEAR85                )86            else:87                print(f"Skipping item {item_idx} (label: {label}) due to empty decoded mask before resize.")88                continue89        except cv2.error as e:90            print(f"Error resizing mask for item {item_idx} (label: {label}): {e}")91            continue92 93        full_mask = np.zeros((img_height, img_width), dtype=np.uint8)  # Start with a black mask (0s)94        try:95            full_mask[abs_y0:abs_y1, abs_x0:abs_x1] = resized_bbox_mask  # Place the white mask in the right location96        except ValueError as e:97            print(f"Error placing mask for item {item_idx} (label: {label}): {e}. Shape mismatch: slice=({bbox_height},{bbox_width}), resized={resized_bbox_mask.shape}. Attempting correction.")98            try:99                resized_bbox_mask_corrected = cv2.resize(bbox_mask, (bbox_width, bbox_height), interpolation=cv2.INTER_LINEAR)100                full_mask[abs_y0:abs_y1, abs_x0:abs_x1] = resized_bbox_mask_corrected101                print("  -> Corrected placement.")102            except Exception as inner_e:103                print(f"  -> Failed to correct placement: {inner_e}")104                continue105 106        segmentation_data.append((full_mask, label))107 108    return segmentation_data109 110 111 112def create_binary_mask_overlay(113    img: PILImage.Image,114    segmentation_data: list[tuple[np.ndarray, str]],115    alpha: float = 0.8116) -> np.ndarray:117 118    binary_mask = np.zeros(img.size[::-1], dtype=np.uint8)119 120    for mask, label in segmentation_data:121        if mask is not None and mask.shape == binary_mask.shape:122            binary_mask = np.maximum(binary_mask, mask)123 124    result = np.zeros_like(binary_mask, dtype=np.uint8)125    result[binary_mask > 0] = 255126 127    return result128 129 130def segment_image(prompt: str, image: Union[str, PILImage.Image]) -> Tuple[np.ndarray, List[tuple[np.ndarray, str]]]:131 132    # --- Normalize inputs ---133    if isinstance(image, str):134        img = PILImage.open(image).convert("RGB")135    else:136        img = image.convert("RGB")137 138    img_width, img_height = img.size139 140    # --- Client setup (no hardcoded secrets) ---141    api_key = os.getenv("GEMINI_API_KEY") 142    if not api_key:143        raise RuntimeError("GEMINI_API_KEY environment variable not set.")144    client = genai.Client(api_key=api_key)145 146    # --- Model + config (kept simple and local to function) ---147    model_id = "gemini-2.5-flash"148    safety_settings = [149        types.SafetySetting(150            category="HARM_CATEGORY_DANGEROUS_CONTENT",151            threshold="BLOCK_ONLY_HIGH",152        ),153    ]154 155    # --- Inference ---156    # Note: the SDK accepts [prompt, PILImage] in `contents` as in your original code.157    response = client.models.generate_content(158        model=model_id,159        contents=[prompt, img],160        config=types.GenerateContentConfig(161            temperature=0.2,162            safety_settings=safety_settings,163            thinking_config=types.ThinkingConfig(164            thinking_budget=0)165        )166    )167    result_text = response.text or ""168 169    # --- Post-process into masks ---170    segmentation_data = generate_mask(171        result_text,172        img_height=img_height,173        img_width=img_width174    )175 176    # Always return something predictable177    if not segmentation_data:178        # empty binary mask, same HxW179        return (np.zeros((img_height, img_width), dtype=np.uint8), [])180 181    binary_mask = create_binary_mask_overlay(img, segmentation_data, alpha=0.8)182    return binary_mask, segmentation_data