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()