Class IconPictureProcessor

modules/icon_picture_processor.py:229–411

Process icon/picture elements: filter, crop, optional RMBG, base64, XML fragments.

# ======================== Icon/Picture processor ========================
class IconPictureProcessor(BaseProcessor):
    """Process icon/picture elements: filter, crop, optional RMBG, base64, XML fragments."""

    # Types that use RMBG for background removal; others keep original crop
    RMBG_TYPES = {"icon", "logo", "symbol", "emoji", "button", "arrow"}
    

    # Types that keep background (crop only)
    KEEP_BG_TYPES = {
        "picture", "photo", "chart", "function_graph", "screenshot", "image", "diagram",
        "graph", "line graph", "bar graph", "heatmap", "scatter plot", "histogram", "pie chart"
    }
    
    # Max element area ratio (skip if element area > this fraction of image)
    MAX_AREA_RATIO = 0.75

    
    def __init__(
        self,
        config=None,
        rmbg_model_path: str = None,
    ):
        super().__init__(config)
        self._rmbg_model: Optional[RMBGModel] = None
        self._rmbg_model_path = rmbg_model_path

    def load_rmbg_model(self):
        """Load RMBG model."""
        if self._rmbg_model is None:
            self._rmbg_model = RMBGModel(self._rmbg_model_path)
        if not self._rmbg_model.is_loaded:
            self._rmbg_model.load()

    def load_model(self):
        """Alias: load RMBG model."""
        self.load_rmbg_model()

    def process(self, context: ProcessingContext) -> ProcessingResult:
        """Process icon/picture elements in context."""
        self._log("Processing Icon/Picture elements")
        self.load_rmbg_model()

        # Load image
        if not context.image_path or not os.path.exists(context.image_path):
            return ProcessingResult(
                success=False,
                error_message="Invalid image path"
            )
        
        original_image = Image.open(context.image_path).convert("RGB")
        cv2_image = cv2.imread(context.image_path)
        
        # Filter elements to process
        elements_to_process = self._get_elements_to_process(context.elements)
        
        self._log(f"Elements to process: {len(elements_to_process)}")
        
        processed_count = 0