Update app.py

app.py

@@ -1,88 +1,74 @@
-import gradio as gr
-import numpy as np
+# no cpu required
+#TODO: update to gpu usage
+from transformers import pipeline, SamModel, SamProcessor
 import torch
-from PIL import Image, ImageDraw
-import requests
-from transformers import SamModel, SamProcessor
-import cv2
-from typing import List
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# we load model and processor
-model = SamModel.from_pretrained("jadechoghari/robustsam-vit-base").to(device)
-processor = SamProcessor.from_pretrained("jadechoghari/robustsam-vit-base")
+import numpy as np
+import spaces
 
-cache_data = None
+checkpoint = "google/owlv2-base-patch16-ensemble"
+detector = pipeline(model=checkpoint, task="zero-shot-object-detection")
+sam_model = SamModel.from_pretrained("jadechoghari/robustsam-vit-base")
+sam_processor = SamProcessor.from_pretrained("jadechoghari/robustsam-vit-base")
 
-def mask_2_dots(mask: np.ndarray) -> List[List[int]]:
-    gray = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
-    _, thresh = cv2.threshold(gray, 127, 255, 0)
-    kernel = np.ones((5,5),np.uint8)
-    closed = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
-    contours, _ = cv2.findContours(closed, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
-    points = []
-    for contour in contours:
-        moments = cv2.moments(contour)
-        cx = int(moments['m10']/moments['m00'])
-        cy = int(moments['m01']/moments['m00'])
-        points.append([cx, cy])
-    return [points]
+@spaces.GPU
+def query(image, texts, threshold):
+  texts = texts.split(",")
 
-@torch.no_grad()
-def foward_pass(image_input: np.ndarray, points: List[List[int]]) -> np.ndarray:
-    global cache_data
-    image_input = Image.fromarray(image_input)
-    inputs = processor(image_input, input_points=points, return_tensors="pt").to(device)
-    if not cache_data or not torch.equal(inputs['pixel_values'],cache_data[0]):
-        embedding = model.get_image_embeddings(inputs["pixel_values"])
-        pixels = inputs["pixel_values"]
-        cache_data = [pixels, embedding]
-    del inputs["pixel_values"]
+  predictions = detector(
+    image,
+    candidate_labels=texts,
+    threshold=threshold
+  )
 
-    outputs = model.forward(image_embeddings=cache_data[1], **inputs)
-    masks = processor.image_processor.post_process_masks(
-        outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
-    )
-    masks = masks[0].squeeze(0).numpy().transpose(1, 2, 0)
+  result_labels = []
+  for pred in predictions:
 
-    return masks
+    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)]
 
-def main_func(inputs) -> List[Image.Image]:
-    dots = inputs['mask']
-    points = mask_2_dots(dots)
-    image_input = inputs['image']
-    masks = foward_pass(image_input, points)
+    inputs = sam_processor(
+            image,
+            input_boxes=[[[box]]],
+            return_tensors="pt"
+        )
 
-    image_input = Image.fromarray(image_input)
-    draw = ImageDraw.Draw(image_input)
-    for point in points[0]:
-        draw.ellipse((point[0] - 10, point[1] - 10, point[0] + 10, point[1] + 10), fill="red")
+    with torch.no_grad():
+        outputs = sam_model(**inputs)
 
-    pred_masks = [image_input]
-    for i in range(masks.shape[2]):
-        pred_masks.append(Image.fromarray((masks[:,:,i] * 255).astype(np.uint8)))
+    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, ...]
 
-    return pred_masks
+    from PIL import Image, ImageDraw
+    # Convert mask to image format and overlay on the original image
+    mask_image = Image.fromarray((mask[0] * 255).astype(np.uint8))
+    mask_image = mask_image.convert("L")  # Convert to grayscale for transparency
+    mask_image = mask_image.resize(image.size)
 
-def reset_data():
-    global cache_data
-    cache_data = None
+    # Create an alpha mask for transparency
+    alpha_mask = Image.new("L", mask_image.size, 128)  # Adjust transparency level here
+    image.paste(mask_image, (0, 0), alpha_mask)  # Overlay the mask on the image
 
-with gr.Blocks() as demo:
-    gr.Markdown("# How to use")
-    gr.Markdown("To start, input an image, then use the brush to create dots on the object which you want to segment, don't worry if your dots aren't perfect as the code will find the middle of each drawn item. Then press the segment button to create masks for the object that the dots are on.")
-    gr.Markdown("# Demo to run Robust Segment Anything base model")
-    gr.Markdown("""This app uses the [Robust Segment Anything](https://huggingface.co/jadechoghari/robustsam-vit-base) model from Snap Research to get a mask from a points in an image.
-    """)
-    with gr.Tab("Flip Image"):
-        with gr.Row():
-            image_input = gr.ImageEditor()
-            image_output = gr.Gallery()
-        
-        image_button = gr.Button("Segment Image")
+    # Save the annotated image
+    image.save("annotated_image.png")
+    print("saved image")
+    result_labels.append((mask, label))
+  return image, result_labels
 
-    image_button.click(main_func, inputs=image_input, outputs=image_output)
-    image_input.upload(reset_data)
+import gradio as gr
 
-demo.launch()
+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. 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="OWL 🤝 SAM",
+    description=description
+)
+demo.launch(debug=True)