Create app.py

app.py

@@ -0,0 +1,84 @@
+import torch
+import gradio as gr
+from transformers import Owlv2Processor, Owlv2ForObjectDetection
+import spaces
+
+# Use GPU if available
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
+
+# Load the OWLv2 model and processor
+model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble").to(device)
+processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemble")
+
+# Define default text queries relevant to home interior & renovation defects.
+default_queries = (
+    "pipe defect, rust on pipe, cracked pipe, plastic pipe defect, metal pipe defect, "
+    "water damage on wall, mold on wall, broken sink, damaged cabinet, faulty door"
+)
+
+@spaces.GPU
+def query_image(img, text_queries, score_threshold):
+    # Use default queries if none provided
+    if not text_queries.strip():
+        text_queries = default_queries
+    # Split and clean text queries into a list
+    queries = [q.strip() for q in text_queries.split(",") if q.strip()]
+    
+    # Determine target size based on the image dimensions
+    size = max(img.shape[:2])
+    target_sizes = torch.Tensor([[size, size]])
+    
+    # Process inputs
+    inputs = processor(text=queries, images=img, return_tensors="pt").to(device)
+    
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    # Bring outputs to CPU and post-process them
+    outputs.logits = outputs.logits.cpu()
+    outputs.pred_boxes = outputs.pred_boxes.cpu()
+    results = processor.post_process_object_detection(outputs=outputs, target_sizes=target_sizes)
+    boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["labels"]
+    
+    result_labels = []
+    for box, score, label in zip(boxes, scores, labels):
+        if score < score_threshold:
+            continue
+        # OWLv2 returns label indices corresponding to the order of the input queries.
+        if label.item() < len(queries):
+            result_label = queries[label.item()]
+        else:
+            result_label = "unknown"
+        box = [int(i) for i in box.tolist()]
+        result_labels.append((box, result_label))
+    
+    return img, result_labels
+
+description = """
+This demo uses OWLv2 for zero-shot object detection, specifically tailored for home interior and renovation defects.
+Enter comma-separated text queries describing issues relevant to home renovations—for example:
+"pipe defect, rust on pipe, cracked pipe, plastic pipe defect, metal pipe defect, water damage on wall, mold on wall, broken sink, damaged cabinet, faulty door".
+If left blank, a default set of queries will be used.
+"""
+
+demo = gr.Interface(
+    fn=query_image,
+    inputs=[
+        gr.Image(type="pil", label="Upload an Image"),
+        gr.Textbox(default=default_queries, label="Text Queries"),
+        gr.Slider(0, 1, value=0.1, label="Score Threshold")
+    ],
+    outputs=[gr.Image(label="Annotated Image"), "json"],
+    title="Zero-Shot Home Renovation Defect Detection with OWLv2",
+    description=description,
+    examples=[
+        # Replace these example paths with your sample images if available.
+        ["assets/pipe_sample.jpg", default_queries, 0.11],
+        ["assets/kitchen_renovation.jpg", default_queries, 0.1],
+    ],
+)
+
+demo.launch()