Spaces:
Sleeping
Sleeping
| """ | |
| app.py — OWLv2 / SAM2 image labeling UI | |
| Tab 1 — Test: upload one image, pick Detection or Segmentation mode, | |
| tune prompts/threshold/size, see instant annotated results. | |
| Tab 2 — Batch: upload multiple images, run in the chosen mode, download a ZIP | |
| containing resized images + coco_export.json. | |
| All artifacts live in a system temp directory — nothing is written to the project. | |
| """ | |
| from __future__ import annotations | |
| # spaces MUST be imported before torch initialises CUDA (i.e. before any | |
| # autolabel import). Do this first, before everything else. | |
| try: | |
| import spaces as _spaces # type: ignore | |
| _ZERO_GPU = True | |
| except (ImportError, RuntimeError): | |
| _spaces = None | |
| _ZERO_GPU = False | |
| import os | |
| os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") | |
| import logging | |
| import shutil | |
| import tempfile | |
| import zipfile | |
| from pathlib import Path | |
| from typing import Optional | |
| import gradio as gr | |
| import numpy as np | |
| from dotenv import load_dotenv | |
| from PIL import Image, ImageDraw, ImageFont | |
| load_dotenv() | |
| from autolabel.config import settings | |
| from autolabel.detect import infer as _owlv2_infer | |
| from autolabel.export import build_coco | |
| from autolabel.segment import load_sam2, segment_with_boxes | |
| from autolabel.utils import save_json, setup_logging | |
| setup_logging(logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| # Temp directory for this session — cleaned up by the OS on reboot | |
| _TMPDIR = Path(tempfile.mkdtemp(prefix="autolabel_")) | |
| logger.info("Session temp dir: %s", _TMPDIR) | |
| # --------------------------------------------------------------------------- | |
| # Image sizing | |
| # --------------------------------------------------------------------------- | |
| _SIZE_OPTIONS = { | |
| "As is": None, | |
| "416 × 416": (416, 416), | |
| "480 × 480": (480, 480), | |
| "512 × 512": (512, 512), | |
| "640 × 640": (640, 640), | |
| "736 × 736": (736, 736), | |
| "896 × 896": (896, 896), | |
| "1024 × 1024": (1024, 1024), | |
| } | |
| _SIZE_LABELS = list(_SIZE_OPTIONS.keys()) | |
| def _resize(pil: Image.Image, size_label: str) -> Image.Image: | |
| target = _SIZE_OPTIONS[size_label] | |
| if target is None: | |
| return pil | |
| return pil.resize(target, Image.LANCZOS) | |
| # --------------------------------------------------------------------------- | |
| # Colours & annotation | |
| # --------------------------------------------------------------------------- | |
| _PALETTE = [ | |
| (52, 211, 153), (251, 146, 60), (96, 165, 250), (248, 113, 113), | |
| (167, 139, 250),(250, 204, 21), (34, 211, 238), (244, 114, 182), | |
| (74, 222, 128), (232, 121, 249), (125, 211, 252), (253, 186, 116), | |
| (110, 231, 183),(196, 181, 253), (253, 164, 175), (134, 239, 172), | |
| ] | |
| def _colour_for(label: str, prompts: list[str]) -> tuple[int, int, int]: | |
| try: | |
| return _PALETTE[prompts.index(label) % len(_PALETTE)] | |
| except ValueError: | |
| return _PALETTE[hash(label) % len(_PALETTE)] | |
| def _annotate( | |
| pil_image: Image.Image, | |
| detections: list[dict], | |
| prompts: list[str], | |
| mode: str = "Detection", | |
| ) -> Image.Image: | |
| """Draw bounding boxes (+ mask overlays in Segmentation mode) on *pil_image*.""" | |
| img = pil_image.copy().convert("RGBA") | |
| # --- Segmentation: paint semi-transparent mask overlays first --- | |
| if mode == "Segmentation": | |
| overlay = Image.new("RGBA", img.size, (0, 0, 0, 0)) | |
| for det in detections: | |
| mask = det.get("mask") | |
| if mask is None or not isinstance(mask, np.ndarray): | |
| continue | |
| r, g, b = _colour_for(det["label"], prompts) | |
| mask_rgba = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8) | |
| mask_rgba[mask] = [r, g, b, 100] # semi-transparent fill | |
| overlay = Image.alpha_composite(overlay, Image.fromarray(mask_rgba, "RGBA")) | |
| img = Image.alpha_composite(img, overlay) | |
| # --- Bounding boxes and labels (both modes) --- | |
| draw = ImageDraw.Draw(img, "RGBA") | |
| try: | |
| font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", size=18) | |
| except Exception: | |
| font = ImageFont.load_default() | |
| for det in detections: | |
| x1, y1, x2, y2 = det["box_xyxy"] | |
| r, g, b = _colour_for(det["label"], prompts) | |
| draw.rectangle([x1, y1, x2, y2], outline=(r, g, b), width=3) | |
| tag = f"{det['label']} {det['score']:.2f}" | |
| bbox = draw.textbbox((x1, y1), tag, font=font) | |
| draw.rectangle([bbox[0]-3, bbox[1]-3, bbox[2]+3, bbox[3]+3], fill=(r, g, b, 210)) | |
| draw.text((x1, y1), tag, fill=(255, 255, 255), font=font) | |
| return img.convert("RGB") | |
| # --------------------------------------------------------------------------- | |
| # OWLv2 model (cached) | |
| # --------------------------------------------------------------------------- | |
| _owlv2_cache: dict = {} | |
| def _get_owlv2(): | |
| if settings.model not in _owlv2_cache: | |
| _owlv2_cache.clear() | |
| from transformers import Owlv2ForObjectDetection, Owlv2Processor | |
| logger.info("Loading OWLv2 %s on %s …", settings.model, settings.device) | |
| processor = Owlv2Processor.from_pretrained(settings.model) | |
| model = Owlv2ForObjectDetection.from_pretrained( | |
| settings.model, torch_dtype=settings.torch_dtype | |
| ).to(settings.device) | |
| model.eval() | |
| _owlv2_cache[settings.model] = (processor, model) | |
| logger.info("OWLv2 ready.") | |
| return _owlv2_cache[settings.model] | |
| # --------------------------------------------------------------------------- | |
| # SAM2 model (cached) | |
| # --------------------------------------------------------------------------- | |
| _sam2_cache: dict = {} | |
| _SAM2_MODEL_ID = "facebook/sam2-hiera-tiny" | |
| def _get_sam2(): | |
| if _SAM2_MODEL_ID not in _sam2_cache: | |
| processor, model = load_sam2(settings.device, _SAM2_MODEL_ID) | |
| _sam2_cache[_SAM2_MODEL_ID] = (processor, model) | |
| return _sam2_cache[_SAM2_MODEL_ID] | |
| # --------------------------------------------------------------------------- | |
| # Shared inference helpers | |
| # --------------------------------------------------------------------------- | |
| def _infer_on_device( | |
| pil_image: Image.Image, | |
| prompts: list[str], | |
| threshold: float, | |
| mode: str, | |
| device: str, | |
| dtype, | |
| ) -> list[dict]: | |
| """Run OWLv2 (+ optional SAM2) with explicit device/dtype. | |
| In ZeroGPU mode this is called inside the @spaces.GPU context so CUDA is | |
| available; locally it uses whatever settings.device resolved to. | |
| """ | |
| processor, owlv2 = _get_owlv2() | |
| owlv2.to(device) | |
| try: | |
| detections = _owlv2_infer( | |
| pil_image, processor, owlv2, prompts, threshold, device, dtype, | |
| ) | |
| finally: | |
| if _ZERO_GPU: | |
| owlv2.to("cpu") # release VRAM back to ZeroGPU pool | |
| if mode == "Segmentation" and detections: | |
| sam2_processor, sam2_model = _get_sam2() | |
| sam2_model.to(device) | |
| try: | |
| detections = segment_with_boxes( | |
| pil_image, detections, sam2_processor, sam2_model, device | |
| ) | |
| finally: | |
| if _ZERO_GPU: | |
| sam2_model.to("cpu") | |
| return detections | |
| if _ZERO_GPU: | |
| def _run_detection( | |
| pil_image: Image.Image, | |
| prompts: list[str], | |
| threshold: float, | |
| mode: str, | |
| ) -> list[dict]: | |
| """ZeroGPU entry-point: GPU is allocated for the duration of this call.""" | |
| import torch | |
| return _infer_on_device( | |
| pil_image, prompts, threshold, mode, | |
| device="cuda", dtype=torch.float16, | |
| ) | |
| else: | |
| def _run_detection( | |
| pil_image: Image.Image, | |
| prompts: list[str], | |
| threshold: float, | |
| mode: str, | |
| ) -> list[dict]: | |
| return _infer_on_device( | |
| pil_image, prompts, threshold, mode, | |
| device=settings.device, dtype=settings.torch_dtype, | |
| ) | |
| def _parse_prompts(text: str) -> list[str]: | |
| return [p.strip() for p in text.split(",") if p.strip()] | |
| # --------------------------------------------------------------------------- | |
| # Object crops | |
| # --------------------------------------------------------------------------- | |
| def _make_crops( | |
| pil_image: Image.Image, | |
| detections: list[dict], | |
| prompts: list[str], | |
| mode: str, | |
| ) -> list[tuple[Image.Image, str]]: | |
| """Return one (cropped PIL image, caption) pair per detection. | |
| Detection mode: plain bounding-box crop with a coloured border. | |
| Segmentation mode: tight crop around the mask's nonzero region; pixels | |
| outside the mask are set to white for a clean cutout. | |
| """ | |
| crops: list[tuple[Image.Image, str]] = [] | |
| img_w, img_h = pil_image.size | |
| for det in detections: | |
| x1, y1, x2, y2 = det["box_xyxy"] | |
| x1 = max(0, int(x1)); y1 = max(0, int(y1)) | |
| x2 = min(img_w, int(x2)); y2 = min(img_h, int(y2)) | |
| if x2 <= x1 or y2 <= y1: | |
| continue | |
| r, g, b = _colour_for(det["label"], prompts) | |
| if mode == "Segmentation": | |
| mask = det.get("mask") | |
| if mask is not None and isinstance(mask, np.ndarray): | |
| # Find the tight bounding box of the mask's nonzero region | |
| rows = np.any(mask, axis=1) | |
| cols = np.any(mask, axis=0) | |
| if rows.any() and cols.any(): | |
| r_min, r_max = int(np.where(rows)[0][0]), int(np.where(rows)[0][-1]) | |
| c_min, c_max = int(np.where(cols)[0][0]), int(np.where(cols)[0][-1]) | |
| mask_tight = mask[r_min:r_max + 1, c_min:c_max + 1] | |
| region = np.array( | |
| pil_image.crop((c_min, r_min, c_max + 1, r_max + 1)).convert("RGB") | |
| ) | |
| # White background outside the mask | |
| region[~mask_tight] = [255, 255, 255] | |
| crop_rgb = Image.fromarray(region) | |
| else: | |
| crop_rgb = pil_image.crop((x1, y1, x2, y2)).convert("RGB") | |
| else: | |
| crop_rgb = pil_image.crop((x1, y1, x2, y2)).convert("RGB") | |
| else: | |
| crop_rgb = pil_image.crop((x1, y1, x2, y2)).convert("RGB") | |
| # Coloured border | |
| bordered = Image.new("RGB", (crop_rgb.width + 6, crop_rgb.height + 6), (r, g, b)) | |
| bordered.paste(crop_rgb, (3, 3)) | |
| caption = f"{det['label']} {det['score']:.2f}" | |
| crops.append((bordered, caption)) | |
| return crops | |
| # --------------------------------------------------------------------------- | |
| # Tab 1 — Test | |
| # --------------------------------------------------------------------------- | |
| def run_test( | |
| image_np: Optional[np.ndarray], | |
| prompts_text: str, | |
| threshold: float, | |
| size_label: str, | |
| mode: str, | |
| ): | |
| if image_np is None or not prompts_text.strip(): | |
| return image_np, [], [] | |
| prompts = _parse_prompts(prompts_text) | |
| if not prompts: | |
| return image_np, [], [] | |
| pil = _resize(Image.fromarray(image_np), size_label) | |
| detections = _run_detection(pil, prompts, threshold, mode) | |
| table = [ | |
| [i + 1, d["label"], f"{d['score']:.3f}", | |
| f"[{d['box_xyxy'][0]:.0f}, {d['box_xyxy'][1]:.0f}, " | |
| f"{d['box_xyxy'][2]:.0f}, {d['box_xyxy'][3]:.0f}]"] | |
| for i, d in enumerate(detections) | |
| ] | |
| crops = _make_crops(pil, detections, prompts, mode) | |
| return np.array(_annotate(pil, detections, prompts, mode)), table, crops | |
| # --------------------------------------------------------------------------- | |
| # Tab 2 — Batch | |
| # --------------------------------------------------------------------------- | |
| def run_batch(files, prompts_text: str, threshold: float, size_label: str, mode: str): | |
| if not files or not prompts_text.strip(): | |
| return [], "Upload images and enter prompts to get started.", None | |
| prompts = _parse_prompts(prompts_text) | |
| if not prompts: | |
| return [], "No valid prompts.", None | |
| # Fresh temp dir for this run | |
| run_dir = _TMPDIR / "current_run" | |
| if run_dir.exists(): | |
| shutil.rmtree(run_dir) | |
| images_dir = run_dir / "images" | |
| images_dir.mkdir(parents=True) | |
| gallery: list[Image.Image] = [] | |
| total_dets = 0 | |
| for f in files: | |
| try: | |
| src = Path(f.name if hasattr(f, "name") else str(f)) | |
| pil = _resize(Image.open(src).convert("RGB"), size_label) | |
| w, h = pil.size | |
| detections = _run_detection(pil, prompts, threshold, mode) | |
| total_dets += len(detections) | |
| # Save resized image (included in the ZIP) | |
| img_name = src.name | |
| pil.save(images_dir / img_name) | |
| # Per-image JSON consumed by build_coco. | |
| # Drop numpy mask arrays — they are not JSON-serialisable. | |
| json_dets = [ | |
| {k: v for k, v in det.items() if k != "mask"} | |
| for det in detections | |
| ] | |
| save_json( | |
| {"image_path": img_name, "image_width": w, | |
| "image_height": h, "detections": json_dets}, | |
| run_dir / (src.stem + ".json"), | |
| ) | |
| gallery.append(_annotate(pil, detections, prompts, mode)) | |
| except Exception: | |
| logger.exception("Failed to process %s", f) | |
| # Build COCO JSON | |
| coco = build_coco(run_dir) | |
| coco_path = run_dir / "coco_export.json" | |
| if coco: | |
| save_json(coco, coco_path) | |
| # Package everything into a ZIP | |
| zip_path = run_dir / "autolabel_export.zip" | |
| with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf: | |
| if coco_path.exists(): | |
| zf.write(coco_path, "coco_export.json") | |
| for img_file in sorted(images_dir.iterdir()): | |
| zf.write(img_file, f"images/{img_file.name}") | |
| n_ann = len(coco.get("annotations", [])) if coco else 0 | |
| size_note = f" · resized to {size_label}" if size_label != "As is" else "" | |
| mode_note = f" · {mode.lower()}" | |
| stats = ( | |
| f"{len(gallery)} image(s) · {total_dets} detection(s) · " | |
| f"{n_ann} annotations{size_note}{mode_note}" | |
| ) | |
| return gallery, stats, str(zip_path) | |
| # --------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------- | |
| _DEFAULT_PROMPTS = ", ".join(settings.prompts[:8]) | |
| _HOW_IT_WORKS_MD = """\ | |
| ## How it works | |
| | Mode | Models | Output | | |
| |------|--------|--------| | |
| | **Detection** | OWLv2 | Bounding boxes + class labels | | |
| | **Segmentation** | OWLv2 → SAM2 | Bounding boxes + pixel masks + COCO polygons | | |
| **Detection** uses [OWLv2](https://huggingface.co/google/owlv2-large-patch14-finetuned), an | |
| open-vocabulary detector that converts your text prompts directly into bounding boxes — no | |
| fixed class list required. | |
| **Segmentation** uses the **Grounded SAM2** pattern: | |
| 1. **OWLv2** reads your text prompts and produces bounding boxes | |
| 2. **SAM2** (`sam2-hiera-tiny`) takes each box as a spatial prompt and refines it into a | |
| pixel-level mask | |
| SAM2 has no concept of text — it only understands spatial prompts (boxes, points, masks). | |
| OWLv2 acts as the *grounding* step, translating words into coordinates that SAM2 can use. | |
| Both models must run in Segmentation mode; Detection mode skips SAM2 entirely. | |
| """ | |
| with gr.Blocks(title="autolabel") as demo: | |
| gr.Markdown("# autolabel — OWLv2 + SAM2") | |
| with gr.Accordion("ℹ️ How it works", open=False): | |
| gr.Markdown(_HOW_IT_WORKS_MD) | |
| with gr.Tabs(): | |
| # ── Tab 1: Test ────────────────────────────────────────────────── | |
| with gr.Tab("🧪 Test"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| t1_image = gr.Image(label="Image — upload, paste, or pick a sample below", | |
| type="numpy", sources=["upload", "clipboard"]) | |
| t1_mode = gr.Radio( | |
| ["Detection", "Segmentation"], | |
| label="Mode", value="Detection", | |
| info="Detection: OWLv2 → boxes only. " | |
| "Segmentation: OWLv2 → boxes → SAM2 → pixel masks.", | |
| ) | |
| t1_prompts = gr.Textbox(label="Prompts (comma-separated)", | |
| value=_DEFAULT_PROMPTS, lines=2) | |
| t1_threshold = gr.Slider(label="Threshold", minimum=0.01, | |
| maximum=0.9, step=0.01, value=settings.threshold) | |
| t1_size = gr.Dropdown(label="Input size", choices=_SIZE_LABELS, | |
| value="As is") | |
| t1_btn = gr.Button("Detect", variant="primary") | |
| with gr.Column(scale=1): | |
| t1_output = gr.Image(label="Result", type="numpy") | |
| t1_table = gr.Dataframe( | |
| headers=["#", "Label", "Score", "Box (xyxy)"], | |
| row_count=(0, "dynamic"), column_count=(4, "fixed"), | |
| ) | |
| t1_crops = gr.Gallery( | |
| label="Object crops", | |
| columns=4, height=220, | |
| object_fit="contain", show_label=True, | |
| ) | |
| # Sample images — click any thumbnail to load it into the image input | |
| _SAMPLES_DIR = Path(__file__).parent / "samples" | |
| gr.Examples( | |
| label="Sample images (click to load)", | |
| examples=[ | |
| [str(_SAMPLES_DIR / "animals.jpg"), "Detection", | |
| "crown, necklace, ball, animal eye", 0.40, "As is"], | |
| [str(_SAMPLES_DIR / "kitchen.jpg"), "Detection", | |
| "apple, banana, orange, broccoli, carrot, bottle, bowl", 0.40, "As is"], | |
| [str(_SAMPLES_DIR / "dog.jpg"), "Detection", | |
| "dog", 0.40, "As is"], | |
| [str(_SAMPLES_DIR / "cat.jpg"), "Detection", | |
| "cat", 0.40, "As is"] | |
| ], | |
| inputs=[t1_image, t1_mode, t1_prompts, t1_threshold, t1_size], | |
| examples_per_page=5, | |
| cache_examples=False, | |
| ) | |
| t1_btn.click( | |
| run_test, | |
| inputs=[t1_image, t1_prompts, t1_threshold, t1_size, t1_mode], | |
| outputs=[t1_output, t1_table, t1_crops], | |
| ) | |
| t1_prompts.submit( | |
| run_test, | |
| inputs=[t1_image, t1_prompts, t1_threshold, t1_size, t1_mode], | |
| outputs=[t1_output, t1_table, t1_crops], | |
| ) | |
| # ── Tab 2: Batch ───────────────────────────────────────────────── | |
| with gr.Tab("📂 Batch"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| t2_files = gr.File(label="Images", file_count="multiple", | |
| file_types=["image"]) | |
| t2_mode = gr.Radio( | |
| ["Detection", "Segmentation"], | |
| label="Mode", value="Detection", | |
| info="Detection: OWLv2 → boxes only. " | |
| "Segmentation: OWLv2 → boxes → SAM2 → pixel masks.", | |
| ) | |
| t2_prompts = gr.Textbox(label="Prompts (comma-separated)", | |
| value=_DEFAULT_PROMPTS, lines=2) | |
| t2_threshold = gr.Slider(label="Threshold", minimum=0.01, | |
| maximum=0.9, step=0.01, value=settings.threshold) | |
| t2_size = gr.Dropdown(label="Input size", choices=_SIZE_LABELS, | |
| value="640 × 640") | |
| t2_btn = gr.Button("Run", variant="primary") | |
| t2_stats = gr.Textbox(label="Stats", interactive=False) | |
| t2_download = gr.DownloadButton( | |
| label="Download ZIP (images + COCO JSON)", | |
| visible=False, variant="secondary", size="sm", | |
| ) | |
| with gr.Column(scale=2): | |
| t2_gallery = gr.Gallery(label="Results", columns=3, | |
| height="auto", object_fit="contain") | |
| def _run_and_reveal(files, prompts_text, threshold, size_label, mode): | |
| gallery, stats, zip_path = run_batch( | |
| files, prompts_text, threshold, size_label, mode | |
| ) | |
| return gallery, stats, gr.update(value=zip_path, visible=zip_path is not None) | |
| t2_btn.click( | |
| _run_and_reveal, | |
| inputs=[t2_files, t2_prompts, t2_threshold, t2_size, t2_mode], | |
| outputs=[t2_gallery, t2_stats, t2_download], | |
| ) | |
| demo.queue(max_size=5) | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860, | |
| share=False, inbrowser=True, theme=gr.themes.Soft()) | |