Update modules/sam_inference.py

modules/sam_inference.py

@@ -8,6 +8,8 @@ from datetime import datetime
 import numpy as np
 import gradio as gr
 
+from PIL import ImageColor  # <-- We need this to convert hex color to (R,G,B)
+
 from modules.model_downloader import (
     AVAILABLE_MODELS, DEFAULT_MODEL_TYPE,
     is_sam_exist,
@@ -21,7 +23,7 @@ from modules.mask_utils import (
     create_mask_combined_images,
     create_mask_gallery,
     create_mask_pixelized_image,
-    create_solid_color_mask_image
+    # create_solid_color_mask_image  <-- We won't call this anymore
 )
 from modules.video_utils import (get_frames_from_dir, create_video_from_frames, get_video_info, extract_frames,
                                  extract_sound, clean_temp_dir, clean_files_with_extension)
@@ -58,10 +60,6 @@ class SamInference:
                    load_video_predictor: bool = False):
         """
         Load the model from the model directory. If the model is not found, download it from the URL.
-
-        Args:
-            model_type (str): The model type to load.
-            load_video_predictor (bool): Load the video predictor model.
         """
         if model_type is None:
             model_type = DEFAULT_MODEL_TYPE
@@ -70,7 +68,6 @@ class SamInference:
         config_dir, config_name = os.path.split(config_path)
 
         filename, url = AVAILABLE_MODELS[model_type]
-
         model_path = os.path.join(self.model_dir, filename)
 
         if not is_sam_exist(model_dir=self.model_dir, model_type=model_type):
@@ -106,10 +103,6 @@ class SamInference:
                                    model_type: Optional[str] = None):
         """
         Initialize the video inference state for the video predictor.
-
-        Args:
-            vid_input (str): The video frames directory.
-            model_type (str): The model type to load.
         """
         if model_type is None:
             model_type = self.current_model_type
@@ -137,18 +130,8 @@ class SamInference:
                       invert_mask: bool = False,
                       **params) -> List[Dict[str, Any]]:
         """
-        Generate masks with Automatic segmentation. Default hyperparameters are in './configs/default_hparams.yaml.'
-
-        Args:
-            image (np.ndarray): The input image.
-            model_type (str): The model type to load.
-            invert_mask (bool): Invert the mask output - used for background masking.
-            **params: The hyperparameters for the mask generator.
-
-        Returns:
-            List[Dict[str, Any]]: The auto-generated mask data.
+        Generate masks with Automatic segmentation.
         """
-
         if self.model is None or self.current_model_type != model_type:
             self.current_model_type = model_type
             self.load_model(model_type=model_type)
@@ -178,20 +161,6 @@ class SamInference:
                       **params) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
         """
         Predict image with prompt data.
-
-        Args:
-            image (np.ndarray): The input image.
-            model_type (str): The model type to load.
-            box (np.ndarray): The box prompt data.
-            point_coords (np.ndarray): The point coordinates prompt data.
-            point_labels (np.ndarray): The point labels prompt data.
-            invert_mask (bool): Invert the mask output - used for background masking.
-            **params: The hyperparameters for the mask generator.
-
-        Returns:
-            np.ndarray: The predicted masks output in CxHxW format.
-            np.ndarray: Array of scores for each mask.
-            np.ndarray: Array of logits in CxHxW format.
         """
         if self.model is None or self.current_model_type != model_type:
             self.current_model_type = model_type
@@ -223,22 +192,8 @@ class SamInference:
                                 labels: Optional[np.ndarray] = None,
                                 box: Optional[np.ndarray] = None) -> Tuple[int, int, torch.Tensor]:
         """
-        Add prediction to the current video inference state. inference state must be initialized before calling this method.
-
-        Args:
-            frame_idx (int): The frame index of the video.
-            obj_id (int): The object id for the frame.
-            inference_state (Dict): The inference state for the video predictor.
-            points (np.ndarray): The point coordinates prompt data.
-            labels (np.ndarray): The point labels prompt data.
-            box (np.ndarray): The box prompt data.
-
-        Returns:
-            int: The frame index of the corresponding prediction.
-            int: The object id of the corresponding prediction.
-            torch.Tensor: The mask logits output in CxHxW format.
+        Add prediction to the current video inference state.
         """
-
         if (self.video_predictor is None or
                 inference_state is None and self.video_inference_state is None):
             logger.exception("Error while predicting frame from video, load video predictor first")
@@ -264,16 +219,7 @@ class SamInference:
     def propagate_in_video(self,
                            inference_state: Optional[Dict] = None,):
         """
-        Propagate in the video with the tracked predictions for each frame. Currently only supports
-        single frame tracking.
-
-        Args:
-            inference_state (Dict): The inference state for the video predictor. Use self.video_inference_state if None.
-
-        Returns:
-            Dict: The video segments with the image and mask data. It has frame index as each key and each key has
-                "image" and "mask" data. "image" key contains the path of the original image file and "mask" key contains
-                the np.ndarray mask output.
+        Propagate in the video with the tracked predictions for each frame.
         """
         if inference_state is None and self.video_inference_state is None:
             logger.exception("Error while propagating in video, load video predictor first")
@@ -312,19 +258,7 @@ class SamInference:
                               invert_mask: bool = False
                               ):
         """
-        Add filter to the preview image with the prompt data. Specially made for gradio app.
-        It adds prediction tracking to the self.video_inference_state and returns the filtered image.
-
-        Args:
-            image_prompt_input_data (Dict): The image prompt data.
-            filter_mode (str): The filter mode to apply. ["Solid Color", "Pixelize"]
-            frame_idx (int): The frame index of the video.
-            pixel_size (int): The pixel size for the pixelize filter.
-            color_hex (str): The color hex code for the solid color filter.
-            invert_mask (bool): Invert the mask output - used for background masking.
-
-        Returns:
-            np.ndarray: The filtered image output.
+        Add filter to the preview image with the prompt data.
         """
         if self.video_predictor is None or self.video_inference_state is None:
             logger.exception("Error while adding filter to preview, load video predictor first")
@@ -357,8 +291,18 @@ class SamInference:
 
         generated_masks = self.format_to_auto_result(masks)
 
+        # ---------------------------
+        # Modified solid color branch
+        # ---------------------------
         if filter_mode == COLOR_FILTER:
-            image = create_solid_color_mask_image(image, generated_masks, color_hex)
+            # Make entire background black, fill the mask area with the chosen color
+            color_rgb = ImageColor.getcolor(color_hex, "RGB")
+            blacked = np.zeros_like(image, dtype=np.uint8)
+            # If there are multiple mask segments, fill them all with the chosen color
+            for m in generated_masks:
+                seg = m["segmentation"]
+                blacked[seg > 0] = color_rgb
+            image = blacked
 
         elif filter_mode == PIXELIZE_FILTER:
             image = create_mask_pixelized_image(image, generated_masks, pixel_size)
@@ -374,22 +318,8 @@ class SamInference:
                               invert_mask: bool = False
                               ):
         """
-        Create a whole filtered video with video_inference_state. Currently only one frame tracking is supported.
-        This needs FFmpeg to run. Returns two output path because of the gradio app.
-
-        Args:
-            image_prompt_input_data (Dict): The image prompt data.
-            filter_mode (str): The filter mode to apply. ["Solid Color", "Pixelize"]
-            frame_idx (int): The frame index of the video.
-            pixel_size (int): The pixel size for the pixelize filter.
-            color_hex (str): The color hex code for the solid color filter.
-            invert_mask (bool): Invert the mask output - used for background masking.
-
-        Returns:
-            str: The output video path.
-            str: The output video path.
+        Create a whole filtered video. Currently only one-frame tracking is supported.
         """
-
         if self.video_predictor is None or self.video_inference_state is None:
             logger.exception("Error while adding filter to preview, load video predictor first")
             raise RuntimeError("Error while adding filter to preview")
@@ -399,13 +329,13 @@ class SamInference:
                              "Please press the eraser button (on the image prompter) and add your prompts again.")
             logger.error(error_message)
             raise gr.Error(error_message, duration=20)
+
         output_dir = os.path.join(self.output_dir, "filter")
 
         clean_files_with_extension(TEMP_OUT_DIR, IMAGE_FILE_EXT)
         self.video_predictor.reset_state(self.video_inference_state)
 
         prompt_frame_image, prompt = image_prompt_input_data["image"], image_prompt_input_data["points"]
-
         point_labels, point_coords, box = self.handle_prompt_data(prompt)
         obj_id = frame_idx
 
@@ -425,8 +355,16 @@ class SamInference:
                 masks = invert_masks(masks)
             masks = self.format_to_auto_result(masks)
 
+            # ---------------------------
+            # Modified solid color branch
+            # ---------------------------
             if filter_mode == COLOR_FILTER:
-                filtered_image = create_solid_color_mask_image(orig_image, masks, color_hex)
+                color_rgb = ImageColor.getcolor(color_hex, "RGB")
+                blacked = np.zeros_like(orig_image, dtype=np.uint8)
+                for m in masks:
+                    seg = m["segmentation"]
+                    blacked[seg > 0] = color_rgb
+                filtered_image = blacked
 
             elif filter_mode == PIXELIZE_FILTER:
                 filtered_image = create_mask_pixelized_image(orig_image, masks, pixel_size)
@@ -454,24 +392,11 @@ class SamInference:
                      *params):
         """
         Divide the layer with the given prompt data and save psd file.
-
-        Args:
-            image_input (np.ndarray): The input image.
-            image_prompt_input_data (Dict): The image prompt data.
-            input_mode (str): The input mode for the image prompt data. ["Automatic", "Box Prompt"]
-            model_type (str): The model type to load.
-            invert_mask (bool): Invert the mask output.
-            *params: The hyperparameters for the mask generator.
-
-        Returns:
-            List[np.ndarray]: List of images by predicted masks.
-            str: The output path of the psd file.
         """
-
         timestamp = datetime.now().strftime("%m%d%H%M%S")
         output_file_name = f"result-{timestamp}.psd"
         output_path = os.path.join(self.output_dir, "psd", output_file_name)
-        # Pre-processed gradio components
+
         hparams = {
             'points_per_side': int(params[0]),
             'points_per_batch': int(params[1]),
@@ -488,14 +413,12 @@ class SamInference:
 
         if input_mode == AUTOMATIC_MODE:
             image = image_input
-
             generated_masks = self.generate_mask(
                 image=image,
                 model_type=model_type,
                 invert_mask=invert_mask,
                 **hparams
             )
-
         elif input_mode == BOX_PROMPT_MODE:
             image = image_prompt_input_data["image"]
             image = np.array(image.convert("RGB"))
@@ -540,14 +463,6 @@ class SamInference:
     ):
         """
         Handle data from ImageInputPrompter.
-
-        Args:
-            prompt_data (Dict): A dictionary containing the 'prompt' key with a list of prompts.
-
-        Returns:
-            point_labels (List): list of points labels.
-            point_coords (List): list of points coords.
-            box (List): list of box datas.
         """
         point_labels, point_coords, box = [], [], []
 
@@ -563,4 +478,4 @@ class SamInference:
         point_coords = np.array(point_coords) if point_coords else None
         box = np.array(box) if box else None
 
-        return point_labels, point_coords, box
+        return point_labels, point_coords, box