finish sam2

app.py

@@ -7,6 +7,8 @@ from transformers import AutoModelForImageSegmentation
 from diffusers import FluxFillPipeline
 from PIL import Image, ImageOps
 from sam2.sam2_image_predictor import SAM2ImagePredictor
+import numpy as np
+import matplotlib.pyplot as plt
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
 
@@ -122,9 +124,29 @@ def rmbg(image=None, url=None):
     return image
 
 
-def mask_generation(image=None, json=None):
+def mask_generation(image=None, d=None):
     predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-tiny")
-    return None
+    predictor.set_image(image)
+    input_point = np.array(d["input_points"])
+    input_label = np.array(d["input_labels"])
+    masks, scores, logits = predictor.predict(
+        point_coords=input_point,
+        point_labels=input_label,
+        multimask_output=True,
+    )
+    sorted_ind = np.argsort(scores)[::-1]
+    masks = masks[sorted_ind]
+    scores = scores[sorted_ind]
+    logits = logits[sorted_ind]
+
+    out = []
+    image = Image.fromarray(image)
+    for i in range(len(masks)):
+        m = Image.fromarray(masks[i] * 255).convert("L")
+        comp = Image.composite(Image.fromarray(image), m, m)
+        out.append((comp, f"image {i}"))
+
+    return out
 
 
 @spaces.GPU
@@ -200,7 +222,13 @@ sam2_tab = gr.Interface(
         gr.Image("image", type="pil"),
         gr.JSON(),
     ],
-    outputs=["image"],
+    outputs=gr.Gallery(),
+    examples=[
+        [
+            "./assets/truck.jpg",
+            {"input_points": [[500, 375], [1125, 625]], "input_labels": [1, 0]},
+        ]
+    ],
 )
 
 demo = gr.TabbedInterface(