Create app.py

app.py

@@ -0,0 +1,63 @@
+from huggingface_hub import hf_hub_download
+
+from transformers import pipeline
+import torch
+import numpy as np
+from sam2.build_sam import build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+import gradio as gr
+
+hf_hub_download(repo_id = "merve/sam2-hiera-tiny", filename="sam2_hiera_tiny.pt", local_dir = "./")
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+CHECKPOINT = f"./sam2_hiera_tiny.pt"
+CONFIG = "sam2_hiera_t.yaml"
+
+sam2_model = build_sam2(CONFIG, CHECKPOINT, device=DEVICE, apply_postprocessing=False)
+predictor = SAM2ImagePredictor(sam2_model)
+
+checkpoint = "google/owlv2-base-patch16-ensemble"
+detector = pipeline(model=checkpoint, task="zero-shot-object-detection", device="cuda")
+
+
+
+def query(image, texts, threshold):
+  texts = texts.split(",")
+
+  predictions = detector(
+    image,
+    candidate_labels=texts,
+    threshold=threshold
+  )
+
+  result_labels = []
+  for pred in predictions:
+    
+    box = pred["box"]
+    score = pred["score"]
+    label = pred["label"]
+    box = [round(pred["box"]["xmin"], 2), round(pred["box"]["ymin"], 2), 
+          round(pred["box"]["xmin"], 2)+round(pred["box"]["xmax"], 2), 
+          round(pred["box"]["ymin"], 2)+round(pred["box"]["ymax"], 2)]
+    predictor.set_image(image)
+
+    mask, scores, logits = predictor.predict(box=box,
+    multimask_output=False)
+    mask = mask[np.newaxis, ...]
+    result_labels.append((mask, label))
+  return image, result_labels
+
+
+description = "This Space combines OWLv2, the state-of-the-art zero-shot object detection model with SAM2, the state-of-the-art mask generation model. SAM2 normally doesn't accept text input. Combining SAM with OWLv2 makes SAM2 text promptable. Try the example or input an image and comma separated candidate labels to segment."
+demo = gr.Interface(
+    query,
+    inputs=[gr.Image(type="pil", label="Image Input"), gr.Textbox(label = "Candidate Labels"), gr.Slider(0, 1, value=0.05, label="Confidence Threshold")],
+    outputs="annotatedimage",
+    title="OWLv2 🤝 SAMv2",
+    description=description,
+    examples=[
+        ["./cats.png", "cat", 0.1],
+    ],
+    cache_examples=True
+)
+demo.launch(debug=True)