File size: 3,263 Bytes
eafc253 f5b1d16 eafc253 |
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 |
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(value=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()
|