room-object-segmentation / app_simple.py
krishnanandiraju's picture
Deploy room segmentation app with caching
d581790 verified
Raw
History Blame Contribute Delete
2.63 kB
"""
Minimal ZeroGPU test - Grounding DINO only
"""
import gradio as gr
import torch
from PIL import Image
from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
import spaces
GDINO_ID = "IDEA-Research/grounding-dino-tiny"
# Load processor only (lightweight)
processor = AutoProcessor.from_pretrained(GDINO_ID)
model = None # Load inside GPU
@spaces.GPU(duration=15)
def detect(image, text_prompt, box_threshold=0.35, text_threshold=0.25):
"""Detect objects using Grounding DINO on ZeroGPU."""
global model
# Load model inside GPU context
if model is None:
print("Loading Grounding DINO...")
model = AutoModelForZeroShotObjectDetection.from_pretrained(GDINO_ID)
print("โœ“ Loaded")
try:
# Process image
print(f"Processing prompt: {text_prompt}")
pil_image = Image.open(image).convert("RGB")
# Run detection
inputs = processor(images=pil_image, text=text_prompt, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
results = processor.post_process_grounded_object_detection(
outputs,
inputs.input_ids,
box_threshold=box_threshold,
text_threshold=text_threshold,
target_sizes=[pil_image.size[::-1]]
)[0]
boxes = results["boxes"].cpu().numpy()
labels = results["labels"]
scores = results["scores"].cpu().numpy()
detections = []
for i in range(len(boxes)):
detections.append({
"label": labels[i],
"score": float(scores[i]),
"box": boxes[i].tolist()
})
return {
"num_detections": len(detections),
"detections": detections
}
except Exception as e:
import traceback
return {"error": f"{type(e).__name__}: {str(e)}\n{traceback.format_exc()}"}
demo = gr.Interface(
fn=detect,
inputs=[
gr.Image(type="filepath", label="Upload image"),
gr.Textbox(value="chair . table . sofa", label="Objects to detect (dot-separated)"),
gr.Slider(0, 1, value=0.35, label="Box threshold"),
gr.Slider(0, 1, value=0.25, label="Text threshold")
],
outputs=gr.JSON(label="Results"),
title="๐Ÿ” Grounding DINO Test",
api_name="detect",
show_error=True
)
if __name__ == "__main__":
demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)