modules/sam_inference.py

from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
from sam2.build_sam import build_sam2, build_sam2_video_predictor
from sam2.sam2_image_predictor import SAM2ImagePredictor
from typing import Dict, List, Optional, Tuple, Any
import torch
import os
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,
    download_sam_model_url
)
from modules.paths import (MODELS_DIR, TEMP_OUT_DIR, TEMP_DIR, MODEL_CONFIGS, OUTPUT_DIR)
from modules.constants import (BOX_PROMPT_MODE, AUTOMATIC_MODE, COLOR_FILTER, PIXELIZE_FILTER, IMAGE_FILE_EXT)
from modules.mask_utils import (
    invert_masks,
    save_psd_with_masks,
    create_mask_combined_images,
    create_mask_gallery,
    create_mask_pixelized_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)
from modules.utils import save_image
from modules.logger_util import get_logger

logger = get_logger()


class SamInference:
    def __init__(self,
                 model_dir: str = MODELS_DIR,
                 output_dir: str = OUTPUT_DIR
                 ):
        self.model = None
        self.available_models = list(AVAILABLE_MODELS.keys())
        self.current_model_type = DEFAULT_MODEL_TYPE
        self.model_dir = model_dir
        self.output_dir = output_dir
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.dtype = torch.float16 if torch.cuda.is_available() else torch.bfloat16
        self.mask_generator = None
        self.image_predictor = None
        self.video_predictor = None
        self.video_inference_state = None
        self.video_info = None

        # Initialization on huggingface space
        for model_name in self.available_models:
            self.load_model(model_type=model_name)

    def load_model(self,
                   model_type: Optional[str] = None,
                   load_video_predictor: bool = False):
        """
        Load the model from the model directory. If the model is not found, download it from the URL.
        """
        if model_type is None:
            model_type = DEFAULT_MODEL_TYPE

        config_path = MODEL_CONFIGS[model_type]
        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):
            logger.info(f"No SAM2 model found, downloading {model_type} model...")
            download_sam_model_url(model_dir=self.model_dir, model_type=model_type)
        logger.info(f"Applying configs to {model_type} model..")

        if load_video_predictor:
            try:
                self.model = None
                self.video_predictor = build_sam2_video_predictor(
                    config_file=config_name,
                    ckpt_path=model_path,
                    device=self.device
                )
                return
            except Exception as e:
                logger.exception("Error while loading SAM2 model for video predictor")

        try:
            self.video_predictor = None
            self.model = build_sam2(
                config_file=config_name,
                ckpt_path=model_path,
                device=self.device
            )
        except Exception as e:
            logger.exception("Error while loading SAM2 model")
            raise RuntimeError(f"Failed to load model") from e

    def init_video_inference_state(self,
                                   vid_input: str,
                                   model_type: Optional[str] = None):
        """
        Initialize the video inference state for the video predictor.
        """
        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
            self.load_model(model_type=model_type, load_video_predictor=True)

        self.video_info = get_video_info(vid_input)
        frames_temp_dir = TEMP_DIR
        clean_temp_dir(frames_temp_dir)
        extract_frames(vid_input, frames_temp_dir)
        if self.video_info.has_sound:
            extract_sound(vid_input, frames_temp_dir)

        if self.video_inference_state is not None:
            self.video_predictor.reset_state(self.video_inference_state)
            self.video_inference_state = None

        self.video_inference_state = self.video_predictor.init_state(video_path=frames_temp_dir)

    def generate_mask(self,
                      image: np.ndarray,
                      model_type: str,
                      invert_mask: bool = False,
                      **params) -> List[Dict[str, Any]]:
        """
        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)
        self.mask_generator = SAM2AutomaticMaskGenerator(
            model=self.model,
            **params
        )
        try:
            generated_masks = self.mask_generator.generate(image)
        except Exception as e:
            logger.exception(f"Error while auto generating masks : {e}")
            raise RuntimeError(f"Failed to generate masks") from e

        if invert_mask:
            generated_masks = [{'segmentation': invert_masks(mask['segmentation']),
                                'area': mask['area']} for mask in generated_masks]

        return generated_masks

    def predict_image(self,
                      image: np.ndarray,
                      model_type: str,
                      box: Optional[np.ndarray] = None,
                      point_coords: Optional[np.ndarray] = None,
                      point_labels: Optional[np.ndarray] = None,
                      invert_mask: bool = False,
                      **params) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Predict image with prompt data.
        """
        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)
        self.image_predictor = SAM2ImagePredictor(sam_model=self.model)
        self.image_predictor.set_image(image)

        try:
            masks, scores, logits = self.image_predictor.predict(
                box=box,
                point_coords=point_coords,
                point_labels=point_labels,
                multimask_output=params["multimask_output"],
            )
        except Exception as e:
            logger.exception(f"Error while predicting image with prompt: {str(e)}")
            raise RuntimeError(f"Failed to predict image with prompt") from e

        if invert_mask:
            masks = invert_masks(masks)

        return masks, scores, logits

    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) -> Tuple[int, int, torch.Tensor]:
        """
        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")

        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,
                frame_idx=frame_idx,
                obj_id=obj_id,
                points=points,
                labels=labels,
                box=box
            )
        except Exception as e:
            logger.exception(f"Error while predicting frame with prompt: {str(e)}")
            raise RuntimeError(f"Failed to predicting frame with prompt") from e

        return out_frame_idx, out_obj_ids, out_mask_logits

    def propagate_in_video(self,
                           inference_state: Optional[Dict] = None,):
        """
        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")

        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
            )
            images = get_frames_from_dir(vid_dir=TEMP_DIR, as_numpy=True)

            with torch.autocast(device_type=self.device, dtype=self.dtype):
                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
                    }
        except Exception as e:
            logger.exception(f"Error while propagating in video: {str(e)}")
            raise RuntimeError(f"Failed to propagate in video") from e

        return video_segments

    def add_filter_to_preview(self,
                              image_prompt_input_data: Dict,
                              filter_mode: str,
                              frame_idx: int,
                              pixel_size: Optional[int] = None,
                              color_hex: Optional[str] = None,
                              invert_mask: bool = False
                              ):
        """
        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")
            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()
        if invert_mask:
            masks = invert_masks(masks)

        generated_masks = self.format_to_auto_result(masks)

        # ---------------------------
        # Modified solid color branch
        # ---------------------------
        if filter_mode == COLOR_FILTER:
            # 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)

        return image

    def create_filtered_video(self,
                              image_prompt_input_data: Dict,
                              filter_mode: str,
                              frame_idx: int,
                              pixel_size: Optional[int] = None,
                              color_hex: Optional[str] = None,
                              invert_mask: bool = False
                              ):
        """
        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")

        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)

        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

        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"]
            if invert_mask:
                masks = invert_masks(masks)
            masks = self.format_to_auto_result(masks)

            # ---------------------------
            # Modified solid color branch
            # ---------------------------
            if filter_mode == COLOR_FILTER:
                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)

            save_image(image=filtered_image, output_dir=TEMP_OUT_DIR)

        if len(video_segments) == 1:
            out_image = save_image(image=filtered_image, output_dir=output_dir)
            return None, out_image

        out_video = create_video_from_frames(
            frames_dir=TEMP_OUT_DIR,
            frame_rate=self.video_info.frame_rate,
            output_dir=output_dir,
        )

        return out_video, out_video

    def divide_layer(self,
                     image_input: np.ndarray,
                     image_prompt_input_data: Dict,
                     input_mode: str,
                     model_type: str,
                     invert_mask: bool = False,
                     *params):
        """
        Divide the layer with the given prompt data and save 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)

        hparams = {
            'points_per_side': int(params[0]),
            'points_per_batch': int(params[1]),
            'pred_iou_thresh': float(params[2]),
            'stability_score_thresh': float(params[3]),
            'stability_score_offset': float(params[4]),
            'crop_n_layers': int(params[5]),
            'box_nms_thresh': float(params[6]),
            'crop_n_points_downscale_factor': int(params[7]),
            'min_mask_region_area': int(params[8]),
            'use_m2m': bool(params[9]),
            'multimask_output': bool(params[10])
        }

        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"))
            prompt = image_prompt_input_data["points"]
            if len(prompt) == 0:
                return [image], []

            point_labels, point_coords, box = self.handle_prompt_data(prompt)

            predicted_masks, scores, logits = self.predict_image(
                image=image,
                model_type=model_type,
                box=box,
                point_coords=point_coords,
                point_labels=point_labels,
                multimask_output=hparams["multimask_output"],
                invert_mask=invert_mask
            )
            generated_masks = self.format_to_auto_result(predicted_masks)

        save_psd_with_masks(image, generated_masks, output_path)
        mask_combined_image = create_mask_combined_images(image, generated_masks)
        gallery = create_mask_gallery(image, generated_masks)
        gallery = [mask_combined_image] + gallery

        return gallery, output_path

    @staticmethod
    def format_to_auto_result(
        masks: np.ndarray
    ):
        """Format the masks to auto result format for convenience."""
        place_holder = 0
        if len(masks.shape) <= 3:
            masks = np.expand_dims(masks, axis=0)
        result = [{"segmentation": mask[0], "area": place_holder} for mask in masks]
        return result

    @staticmethod
    def handle_prompt_data(
        prompt_data: List
    ):
        """
        Handle data from ImageInputPrompter.
        """
        point_labels, point_coords, box = [], [], []

        for x1, y1, left_click_indicator, x2, y2, point_indicator in prompt_data:
            is_point = point_indicator == 4.0
            if is_point:
                point_labels.append(left_click_indicator)
                point_coords.append([x1, y1])
            else:
                box.append([x1, y1, x2, y2])

        point_labels = np.array(point_labels) if point_labels else None
        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