Spaces:
Running on Zero
Running on Zero
File size: 3,779 Bytes
3298147 fd694bc 479d204 fd694bc 479d204 fd694bc d2bf7f8 479d204 d2bf7f8 479d204 d2bf7f8 479d204 d2bf7f8 479d204 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import spaces
import gradio as gr
import numpy as np
import torch
from transformers import (
AutoModelForZeroShotObjectDetection,
AutoProcessor,
SamModel,
SamProcessor,
)
CUDA_DEVICE = "cuda"
SAM_MODEL_ID = "facebook/sam-vit-base"
GROUNDING_DINO_MODEL_ID = "IDEA-Research/grounding-dino-base"
# ZeroGPU enables CUDA emulation during startup, so models should be moved to
# CUDA at module scope. Importing `spaces` before `torch` is important.
sam_model = SamModel.from_pretrained(SAM_MODEL_ID).to(CUDA_DEVICE).eval()
sam_processor = SamProcessor.from_pretrained(SAM_MODEL_ID)
dino_processor = AutoProcessor.from_pretrained(GROUNDING_DINO_MODEL_ID)
dino_model = (
AutoModelForZeroShotObjectDetection.from_pretrained(GROUNDING_DINO_MODEL_ID)
.to(CUDA_DEVICE)
.eval()
)
def infer_dino(image: np.ndarray, text_queries: list[str], score_threshold: float):
query_text = ". ".join(text_queries) + "."
height, width = image.shape[:2]
inputs = dino_processor(
images=image,
text=query_text,
return_tensors="pt",
).to(CUDA_DEVICE)
with torch.inference_mode():
outputs = dino_model(**inputs)
return dino_processor.post_process_grounded_object_detection(
outputs=outputs,
input_ids=inputs.input_ids,
threshold=score_threshold,
text_threshold=score_threshold,
target_sizes=[(height, width)],
)[0]
@spaces.GPU(duration=120)
def query_image(image: np.ndarray, text_queries: str, dino_threshold: float):
if image is None:
raise gr.Error("Please upload an image.")
queries = [query.strip() for query in text_queries.split(",") if query.strip()]
if not queries:
raise gr.Error("Please enter at least one candidate label.")
detections = infer_dino(image, queries, dino_threshold)
boxes = detections["boxes"].detach().cpu()
labels = detections.get("text_labels", detections.get("labels", []))
if len(boxes) == 0:
return image, []
# Process every detected box in one SAM forward pass.
sam_inputs = sam_processor(
images=image,
input_boxes=[boxes.tolist()],
return_tensors="pt",
).to(CUDA_DEVICE)
with torch.inference_mode():
sam_outputs = sam_model(**sam_inputs, multimask_output=True)
masks = sam_processor.image_processor.post_process_masks(
sam_outputs.pred_masks.detach().cpu(),
sam_inputs["original_sizes"].detach().cpu(),
sam_inputs["reshaped_input_sizes"].detach().cpu(),
)[0]
iou_scores = sam_outputs.iou_scores.detach().cpu()[0]
best_mask_indices = iou_scores.argmax(dim=-1)
annotations = []
for detection_index, label in enumerate(labels):
best_mask_index = int(best_mask_indices[detection_index])
mask = masks[detection_index, best_mask_index].numpy().astype(bool)
annotations.append((mask, str(label)))
return image, annotations
description = """
This Space combines Grounding DINO for open-vocabulary object detection with
SAM for text-prompted image segmentation. Enter comma-separated labels such as
`cat, dog, sidewalk, crosswalk`.
"""
demo = gr.Interface(
fn=query_image,
inputs=[
gr.Image(type="numpy", label="Image Input"),
gr.Textbox(
label="Candidate Labels",
placeholder="cat, dog, sidewalk, crosswalk",
),
gr.Slider(
minimum=0.0,
maximum=1.0,
value=0.25,
step=0.01,
label="Grounding DINO Confidence Threshold",
),
],
outputs=gr.AnnotatedImage(label="Segmentation Result"),
title="Grounding DINO + SAM Zero-shot Segmentation",
description=description,
)
if __name__ == "__main__":
demo.launch()
|