update generate mask method

app.py

@@ -8,10 +8,10 @@ from flax.training.common_utils import shard
 from PIL import Image
 from segment_anything import SamPredictor, sam_model_registry, SamAutomaticMaskGenerator
 from diffusers import (
-    UniPCMultistepScheduler,
     FlaxStableDiffusionControlNetPipeline,
     FlaxControlNetModel,
 )
+from transformers import pipeline
 
 import colorsys
 
@@ -69,7 +69,7 @@ with gr.Blocks() as demo:
         submit = gr.Button("Submit")
         clear = gr.Button("Clear")
 
-    def generate_mask(image, evt: gr.SelectData):
+    def generate_mask(image):
         predictor.set_image(image)
         input_point = np.array([120, 21])
         input_label = np.ones(input_point.shape[0])
@@ -82,26 +82,26 @@ with gr.Blocks() as demo:
         # clear torch cache
         torch.cuda.empty_cache()
         mask = Image.fromarray(mask[0, :, :])
-        segs = mask_generator.generate(image)
-        boolean_masks = [s["segmentation"] for s in segs]
-        finseg = np.zeros(
-            (boolean_masks[0].shape[0], boolean_masks[0].shape[1], 3), dtype=np.uint8
-        )
-        # Loop over the boolean masks and assign a unique color to each class
-        for class_id, boolean_mask in enumerate(boolean_masks):
-            hue = class_id * 1.0 / len(boolean_masks)
-            rgb = tuple(int(i * 255) for i in colorsys.hsv_to_rgb(hue, 1, 1))
-            rgb_mask = np.zeros(
-                (boolean_mask.shape[0], boolean_mask.shape[1], 3), dtype=np.uint8
-            )
-            rgb_mask[:, :, 0] = boolean_mask * rgb[0]
-            rgb_mask[:, :, 1] = boolean_mask * rgb[1]
-            rgb_mask[:, :, 2] = boolean_mask * rgb[2]
-            finseg += rgb_mask
+        # segs = mask_generator.generate(image)
+        # boolean_masks = [s["segmentation"] for s in segs]
+        # finseg = np.zeros(
+        #     (boolean_masks[0].shape[0], boolean_masks[0].shape[1], 3), dtype=np.uint8
+        # )
+        # # Loop over the boolean masks and assign a unique color to each class
+        # for class_id, boolean_mask in enumerate(boolean_masks):
+        #     hue = class_id * 1.0 / len(boolean_masks)
+        #     rgb = tuple(int(i * 255) for i in colorsys.hsv_to_rgb(hue, 1, 1))
+        #     rgb_mask = np.zeros(
+        #         (boolean_mask.shape[0], boolean_mask.shape[1], 3), dtype=np.uint8
+        #     )
+        #     rgb_mask[:, :, 0] = boolean_mask * rgb[0]
+        #     rgb_mask[:, :, 1] = boolean_mask * rgb[1]
+        #     rgb_mask[:, :, 2] = boolean_mask * rgb[2]
+        #     finseg += rgb_mask
 
         torch.cuda.empty_cache()
 
-        return mask, finseg
+        return mask
 
     def infer(
         image, prompts, negative_prompts, num_inference_steps=50, seed=4, num_samples=4