Add video propagation

app.py

@@ -191,6 +191,9 @@ class App:
                     btn_generate_preview.click(fn=self.sam_inf.add_filter_to_preview,
                                                inputs=preview_params,
                                                outputs=[img_preview])
+                    btn_generate.click(fn=self.sam_inf.add_filter_to_video,
+                                       inputs=preview_params,
+                                       outputs=None)
 
         self.demo.queue().launch(inbrowser=True)
 

modules/sam_inference.py

@@ -14,7 +14,7 @@ from modules.model_downloader import (
     is_sam_exist,
     download_sam_model_url
 )
-from modules.paths import SAM2_CONFIGS_DIR, MODELS_DIR
+from modules.paths import SAM2_CONFIGS_DIR, MODELS_DIR, TEMP_OUT_DIR, TEMP_DIR
 from modules.constants import BOX_PROMPT_MODE, AUTOMATIC_MODE, COLOR_FILTER, PIXELIZE_FILTER
 from modules.mask_utils import (
     save_psd_with_masks,
@@ -23,6 +23,8 @@ from modules.mask_utils import (
     create_mask_pixelized_image,
     create_solid_color_mask_image
 )
+from modules.video_utils import get_frames_from_dir
+from modules.utils import save_image
 from modules.logger_util import get_logger
 
 MODEL_CONFIGS = {
@@ -45,7 +47,8 @@ class SamInference:
         self.model_dir = model_dir
         self.output_dir = output_dir
         self.model_path = os.path.join(self.model_dir, AVAILABLE_MODELS[DEFAULT_MODEL_TYPE][0])
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.dtype = torch.float16 if torch.cuda.is_available() else torch.float32
         self.mask_generator = None
         self.image_predictor = None
         self.video_predictor = None
@@ -89,8 +92,10 @@ class SamInference:
             raise f"Error while loading SAM2 model!: {e}"
 
     def init_video_inference_state(self,
-                                   model_type: str,
-                                   vid_input: str):
+                                   vid_input: str,
+                                   model_type: Optional[str] = None):
+        if model_type is None:
+            model_type = self.current_model_type
 
         if self.video_predictor is None or model_type != self.current_model_type:
             self.current_model_type = model_type
@@ -141,21 +146,25 @@ class SamInference:
                 multimask_output=params["multimask_output"],
             )
         except Exception as e:
-            logger.exception("Error while predicting image with prompt")
-            raise f"Error while predicting image with prompt: {str(e)}"
+            logger.exception(f"Error while predicting image with prompt: {str(e)}")
+            raise RuntimeError(f"Error while predicting image with prompt: {str(e)}") from e
         return masks, scores, logits
 
-    def predict_frame(self,
-                      frame_idx: int,
-                      obj_id: int,
-                      inference_state: Dict,
-                      points: Optional[np.ndarray] = None,
-                      labels: Optional[np.ndarray] = None,
-                      box: Optional[np.ndarray] = None):
-        if self.video_predictor is None or self.video_inference_state is None:
+    def add_prediction_to_frame(self,
+                                frame_idx: int,
+                                obj_id: int,
+                                inference_state: Optional[Dict] = None,
+                                points: Optional[np.ndarray] = None,
+                                labels: Optional[np.ndarray] = None,
+                                box: Optional[np.ndarray] = None):
+        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")
             raise f"Error while predicting frame from video"
 
+        if inference_state is None:
+            inference_state = self.video_inference_state
+
         try:
             out_frame_idx, out_obj_ids, out_mask_logits = self.video_predictor.add_new_points_or_box(
                 inference_state=inference_state,
@@ -166,15 +175,43 @@ class SamInference:
                 box=box
             )
         except Exception as e:
-            logger.exception("Error while predicting frame with prompt")
-            print(e)
-            raise f"Error while predicting frame with prompt"
+            logger.exception(f"Error while predicting frame with prompt: {str(e)}")
+            raise RuntimeError(f"Failed to predicting frame with prompt: {str(e)}") from e
 
         return out_frame_idx, out_obj_ids, out_mask_logits
 
-    def predict_video(self,
-                      video_input):
-        pass
+    def propagate_in_video(self,
+                           inference_state: Optional[Dict] = None,):
+        if inference_state is None and self.video_inference_state is None:
+            logger.exception("Error while propagating in video, load video predictor first")
+            raise f"Error while propagating in video"
+
+        if inference_state is None:
+            inference_state = self.video_inference_state
+
+        video_segments = {}
+
+        try:
+            generator = self.video_predictor.propagate_in_video(
+                inference_state=inference_state,
+                start_frame_idx=0
+            )
+            cached_images = inference_state["images"]
+            images = get_frames_from_dir(vid_dir=TEMP_DIR, as_numpy=True)
+
+            with torch.autocast(device_type=self.device, dtype=torch.float16):
+                for out_frame_idx, out_obj_ids, out_mask_logits in generator:
+                    mask = (out_mask_logits[0] > 0.0).cpu().numpy()
+                    video_segments[out_frame_idx] = {
+                        "image": images[out_frame_idx],
+                        "mask": mask
+                    }
+                    print("frame_idx: ", out_frame_idx)
+        except Exception as e:
+            logger.exception(f"Error while propagating in video: {str(e)}")
+            raise RuntimeError(f"Failed to propagate in video: {str(e)}") from e
+
+        return video_segments
 
     def add_filter_to_preview(self,
                               image_prompt_input_data: Dict,
@@ -183,49 +220,86 @@ class SamInference:
                               pixel_size: Optional[int] = None,
                               color_hex: Optional[str] = None,
                               ):
-        if not image_prompt_input_data["points"]:
-            error_message = ("Prompt data is empty! Please provide at least one point or box on the image. <br>"
-                             "If you've already added prompts, please press the eraser button "
-                             "and add your prompts again.")
-            logger.error(error_message)
-            raise gr.Error(error_message, duration=20)
-
         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 f"Error while adding filter to preview"
 
+        if not image_prompt_input_data["points"]:
+            error_message = ("No prompt data provided. If this is an incorrect flag, "
+                             "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)
+
         image, prompt = image_prompt_input_data["image"], image_prompt_input_data["points"]
         image = np.array(image.convert("RGB"))
 
         point_labels, point_coords, box = self.handle_prompt_data(prompt)
+        obj_id = frame_idx
+
+        self.video_predictor.reset_state(self.video_inference_state)
+        idx, scores, logits = self.add_prediction_to_frame(
+            frame_idx=frame_idx,
+            obj_id=obj_id,
+            inference_state=self.video_inference_state,
+            points=point_coords,
+            labels=point_labels,
+            box=box
+        )
+        masks = (logits[0] > 0.0).cpu().numpy()
+        generated_masks = self.format_to_auto_result(masks)
 
         if filter_mode == COLOR_FILTER:
-            idx, scores, logits = self.predict_frame(
-                frame_idx=frame_idx,
-                obj_id=0,
-                inference_state=self.video_inference_state,
-                points=point_coords,
-                labels=point_labels,
-                box=box
-            )
-            masks = (logits[0] > 0.0).cpu().numpy()
-            generated_masks = self.format_to_auto_result(masks)
             image = create_solid_color_mask_image(image, generated_masks, color_hex)
 
         elif filter_mode == PIXELIZE_FILTER:
-            idx, scores, logits = self.predict_frame(
-                frame_idx=frame_idx,
-                obj_id=0,
-                inference_state=self.video_inference_state,
-                points=point_coords,
-                labels=point_labels,
-                box=box
-            )
-            masks = (logits[0] > 0.0).cpu().numpy()
-            generated_masks = self.format_to_auto_result(masks)
             image = create_mask_pixelized_image(image, generated_masks, pixel_size)
+
         return image
 
+    def add_filter_to_video(self,
+                            image_prompt_input_data: Dict,
+                            filter_mode: str,
+                            frame_idx: int,
+                            pixel_size: Optional[int] = None,
+                            color_hex: Optional[str] = None,):
+        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 f"Error while adding filter to preview"
+
+        if not image_prompt_input_data["points"]:
+            error_message = ("No prompt data provided. If this is an incorrect flag, "
+                             "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)
+
+        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
+
+        self.video_predictor.reset_state(self.video_inference_state)
+        idx, scores, logits = self.add_prediction_to_frame(
+            frame_idx=frame_idx,
+            obj_id=obj_id,
+            inference_state=self.video_inference_state,
+            points=point_coords,
+            labels=point_labels,
+            box=box
+        )
+
+        video_segments = self.propagate_in_video(inference_state=self.video_inference_state)
+        for frame_index, info in video_segments.items():
+            orig_image, masks = info["image"], info["mask"]
+            masks = self.format_to_auto_result(masks)
+
+            if filter_mode == COLOR_FILTER:
+                filtered_image = create_solid_color_mask_image(orig_image, masks, color_hex)
+
+            elif filter_mode == PIXELIZE_FILTER:
+                filtered_image = create_mask_pixelized_image(orig_image, masks, pixel_size)
+
+            save_image(filtered_image, os.path.join(TEMP_OUT_DIR, "%05d.jpg"))
+
     def divide_layer(self,
                      image_input: np.ndarray,
                      image_prompt_input_data: Dict,