Create app.py

app.py

@@ -0,0 +1,59 @@
+from transformers import pipeline, SamModel, SamProcessor
+import torch
+import numpy as np
+
+checkpoint = "google/owlvit-base-patch16"
+detector = pipeline(model=checkpoint, task="zero-shot-object-detection")
+sam_model = SamModel.from_pretrained("facebook/sam-vit-base").to("cuda")
+sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
+
+def query(image, texts, threshold):
+  texts = texts.split(",")
+  print(texts)
+  print(image.size)
+  predictions = detector(
+    image,
+    candidate_labels=texts,
+  )
+  print(predictions)
+  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"]["xmax"], 2), round(pred["box"]["ymax"], 2)]
+
+    inputs = sam_processor(
+            image,
+            input_boxes=[[[box]]],
+            return_tensors="pt"
+        ).to("cuda")
+
+    with torch.no_grad():
+        outputs = sam_model(**inputs)
+
+    mask = sam_processor.image_processor.post_process_masks(
+        outputs.pred_masks.cpu(),
+        inputs["original_sizes"].cpu(),
+        inputs["reshaped_input_sizes"].cpu()
+    )[0][0][0].numpy()
+    mask = mask[np.newaxis, ...]
+    result_labels.append((mask, label))
+  return image, result_labels
+
+import gradio as gr
+
+description = "This Space combines OWLv2, the state-of-the-art zero-shot object detection model with SAM, the state-of-the-art mask generation model. SAM normally doesn't accept text input. Combining SAM with OWLv2 makes SAM text promptable."
+demo = gr.Interface(
+    query,
+    inputs=[gr.Image(type="pil"), "text", gr.Slider(0, 1, value=0.2)],
+    outputs="annotatedimage",
+    title="OWL 🤝 SAM",
+    #description=description,
+    examples=[
+        ["/content/cats.png", "cat", 0.1],
+    ],
+)
+demo.launch(debug=True)