Support output black or white mask and background

Inference.py

@@ -68,6 +68,9 @@ def parse_args():
     parser.add_argument(
         "--withContours", type=bool, default=False, help="draw the edges of the masks"
     )
+    parser.add_argument(
+        "--bwMask", type=str, default="", help="black or white mask mode, reverse for background and foreground"
+    )
     return parser.parse_args()
 
 
@@ -112,6 +115,7 @@ def main(args):
         point_label = point_label,
         withContours=args.withContours,
         better_quality=args.better_quality,
+        bwMask=args.bwMask
     )
 
 

MORE_USAGES.md

@@ -57,3 +57,27 @@ python Inference.py --model_path ./weights/FastSAM.pt \
                     --withContours True 
 ```
 ![text prompt](assets/more_usages/text_prompt_cat.png)
+
+### use text prompt and output black mask
+Use `--bwMask "black"` to specify the mask color
+```shell
+python Inference.py --model_path ./weights/FastSAM.pt \
+                    --img_path ./images/dogs.jpg \
+                    --text_prompt "the black dog" \
+                    --better_quality True \
+                    --withContours True \
+                    --bwMask "black" 
+```
+![text prompt](assets/more_usages/dogs_black.jpeg)
+
+### use text prompt and output white mask
+Use `--bwMask "white"` to specify the mask color
+```shell
+python Inference.py --model_path ./weights/FastSAM.pt \
+                    --img_path ./images/dogs.jpg \
+                    --text_prompt "the black dog" \
+                    --better_quality True \
+                    --withContours True \
+                    --bwMask "white" 
+```
+![text prompt](assets/more_usages/dogs_white.jpeg)

app_gradio.py

@@ -80,7 +80,11 @@ def segment_everything(
     text="",
     wider=False,
     mask_random_color=True,
+    bw_mask="",
 ):
+    # Reset the default value of bw_mask
+    if bw_mask=="None":
+        bw_mask=""
     input_size = int(input_size)  # 确保 imgsz 是整数
     # Thanks for the suggestion by hysts in HuggingFace.
     w, h = input.size
@@ -111,7 +115,8 @@ def segment_everything(
                        mask_random_color=mask_random_color,
                        bbox=None,
                        use_retina=use_retina,
-                       withContours=withContours,)
+                       withContours=withContours,
+                       bwMask=bw_mask)
     return fig
 
 
@@ -124,6 +129,7 @@ def segment_with_points(
     withContours=True,
     use_retina=True,
     mask_random_color=True,
+    bw_mask="",
 ):
     global global_points
     global global_point_label
@@ -157,7 +163,8 @@ def segment_with_points(
                        mask_random_color=mask_random_color,
                        bbox=None,
                        use_retina=use_retina,
-                       withContours=withContours,)
+                       withContours=withContours,
+                       bwMask=bw_mask)
 
     global_points = []
     global_point_label = []
@@ -342,6 +349,9 @@ def get_points_with_draw(image, label, evt: gr.SelectData):
                         mor_check = gr.Checkbox(value=False, label='better_visual_quality', info='better quality using morphologyEx')
                         retina_check = gr.Checkbox(value=True, label='use_retina', info='draw high-resolution segmentation masks')
                         wider_check = gr.Checkbox(value=False, label='wider', info='wider result')
+                    with gr.Row():
+                        bw_mask = gr.Radio(["None", "black", "white"], value="None", label="Output result image as white or black masks")
+                        rand_color = gr.Checkbox(value=False, label='random', info='mask with random color')
 
                 # Description
                 gr.Markdown(description_e)
@@ -357,6 +367,8 @@ def get_points_with_draw(image, label, evt: gr.SelectData):
                             retina_check,
                             text_box,
                             wider_check,
+                            rand_color,
+                            bw_mask,
                         ],
                         outputs=segm_img_t)
 

fastsam/prompt.py

@@ -100,7 +100,8 @@ def plot_to_result(self,
              mask_random_color=True,
              better_quality=True,
              retina=False,
-             withContours=True) -> np.ndarray:
+             withContours=True,
+             bwMask="") -> np.ndarray:
         if isinstance(annotations[0], dict):
             annotations = [annotation['segmentation'] for annotation in annotations]
         image = self.img
@@ -109,14 +110,17 @@ def plot_to_result(self,
         original_w = image.shape[1]
         if sys.platform == "darwin":
             plt.switch_backend("TkAgg")
-        plt.figure(figsize=(original_w / 100, original_h / 100))
+        bgColor="white"
+        if bwMask == "white":
+            bgColor="black"
+        plt.figure(figsize=(original_w / 100, original_h / 100), facecolor=bgColor, edgecolor=bgColor)
         # Add subplot with no margin.
         plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
         plt.margins(0, 0)
         plt.gca().xaxis.set_major_locator(plt.NullLocator())
         plt.gca().yaxis.set_major_locator(plt.NullLocator())
-
-        plt.imshow(image)
+        if bwMask == "":
+            plt.imshow(image)
         if better_quality:
             if isinstance(annotations[0], torch.Tensor):
                 annotations = np.array(annotations.cpu())
@@ -135,6 +139,7 @@ def plot_to_result(self,
                 retinamask=retina,
                 target_height=original_h,
                 target_width=original_w,
+                bwMask=bwMask,
             )
         else:
             if isinstance(annotations[0], np.ndarray):
@@ -149,6 +154,7 @@ def plot_to_result(self,
                 retinamask=retina,
                 target_height=original_h,
                 target_width=original_w,
+                bwMask=bwMask,
             )
         if isinstance(annotations, torch.Tensor):
             annotations = annotations.cpu().numpy()
@@ -198,7 +204,8 @@ def plot(self,
              mask_random_color=True,
              better_quality=True,
              retina=False,
-             withContours=True):
+             withContours=True,
+             bwMask=""):
         if len(annotations) == 0:
             return None
         result = self.plot_to_result(
@@ -210,6 +217,7 @@ def plot(self,
             better_quality, 
             retina, 
             withContours,
+            bwMask,
         )
 
         path = os.path.dirname(os.path.abspath(output_path))
@@ -230,6 +238,7 @@ def fast_show_mask(
         retinamask=True,
         target_height=960,
         target_width=960,
+        bwMask="",
     ):
         msak_sum = annotation.shape[0]
         height = annotation.shape[1]
@@ -239,12 +248,19 @@ def fast_show_mask(
         sorted_indices = np.argsort(areas)
         annotation = annotation[sorted_indices]
 
+        opacity=0.6
         index = (annotation != 0).argmax(axis=0)
-        if random_color:
+        if bwMask == "white":
+            color = np.ones((msak_sum, 1, 1, 3))
+            opacity=1
+        elif bwMask == "black":
+            color = np.zeros((msak_sum, 1, 1, 3))
+            opacity=1
+        elif random_color:
             color = np.random.random((msak_sum, 1, 1, 3))
         else:
             color = np.ones((msak_sum, 1, 1, 3)) * np.array([30 / 255, 144 / 255, 255 / 255])
-        transparency = np.ones((msak_sum, 1, 1, 1)) * 0.6
+        transparency = np.ones((msak_sum, 1, 1, 1)) * opacity
         visual = np.concatenate([color, transparency], axis=-1)
         mask_image = np.expand_dims(annotation, -1) * visual
 
@@ -287,6 +303,7 @@ def fast_show_mask_gpu(
         retinamask=True,
         target_height=960,
         target_width=960,
+        bwMask="",
     ):
         msak_sum = annotation.shape[0]
         height = annotation.shape[1]
@@ -295,13 +312,20 @@ def fast_show_mask_gpu(
         sorted_indices = torch.argsort(areas, descending=False)
         annotation = annotation[sorted_indices]
         # Find the index of the first non-zero value at each position.
+        opacity=0.6
         index = (annotation != 0).to(torch.long).argmax(dim=0)
-        if random_color:
+        if bwMask == "white":
+            color = torch.ones((msak_sum, 1, 1, 3)).to(annotation.device)
+            opacity=1
+        elif bwMask == "black":
+            color = torch.zeros((msak_sum, 1, 1, 3)).to(annotation.device)
+            opacity=1
+        elif random_color:
             color = torch.rand((msak_sum, 1, 1, 3)).to(annotation.device)
         else:
             color = torch.ones((msak_sum, 1, 1, 3)).to(annotation.device) * torch.tensor([
                 30 / 255, 144 / 255, 255 / 255]).to(annotation.device)
-        transparency = torch.ones((msak_sum, 1, 1, 1)).to(annotation.device) * 0.6
+        transparency = torch.ones((msak_sum, 1, 1, 1)).to(annotation.device) * opacity
         visual = torch.cat([color, transparency], dim=-1)
         mask_image = torch.unsqueeze(annotation, -1) * visual
         # Select data according to the index. The index indicates which batch's data to choose at each position, converting the mask_image into a single batch form.

predict.py

@@ -43,6 +43,10 @@ def predict(
         better_quality: bool = Input(
             description="better quality using morphologyEx", default=False
         ),
+        bw_mask: str = Input(
+            default="",
+            description="empty value to disable, set white or black to get white or black mask"
+        ),
     ) -> Path:
         """Run a single prediction on the model"""
 
@@ -77,6 +81,7 @@ def predict(
             retina=retina,
             text_prompt=text_prompt,
             withContours=withContours,
+            bw_mask=bw_mask
         )
         args.point_prompt = ast.literal_eval(args.point_prompt)
         args.box_prompt = ast.literal_eval(args.box_prompt)
@@ -102,14 +107,18 @@ def predict(
                 args=args,
                 mask_random_color=args.randomcolor,
                 bbox=convert_box_xywh_to_xyxy(args.box_prompt),
+                bwMask=args.bw_mask,
             )
 
         elif args.text_prompt != None:
             results = format_results(results[0], 0)
             annotations = prompt(results, args, text=True)
             annotations = np.array([annotations])
             fast_process(
-                annotations=annotations, args=args, mask_random_color=args.randomcolor
+                annotations=annotations,
+                args=args,
+                mask_random_color=args.randomcolor,
+                bwMask=args.bw_mask,
             )
 
         elif args.point_prompt[0] != [0, 0]:
@@ -122,13 +131,15 @@ def predict(
                 args=args,
                 mask_random_color=args.randomcolor,
                 points=args.point_prompt,
+                bwMask=args.bw_mask,
             )
 
         else:
             fast_process(
                 annotations=results[0].masks.data,
                 args=args,
                 mask_random_color=args.randomcolor,
+                bwMask=args.bw_mask,
             )
 
         out = "/tmp.out.png"

requirements.txt

@@ -12,7 +12,7 @@ tqdm>=4.64.0
 pandas>=1.1.4
 seaborn>=0.11.0
 
-gradio==3.35.2
+gradio==3.41.0
 
 # Ultralytics-----------------------------------
 ultralytics == 8.0.120