Method predict

modules/icon_picture_processor.py:160–217

Background removal; fallback to CPU if GPU fails. Returns RGBA PIL.

(self, image: Image.Image)

        return alpha_resized
    
    def predict(self, image: Image.Image) -> Image.Image:
        """Background removal; fallback to CPU if GPU fails. Returns RGBA PIL."""
        if not self._is_loaded:
            self.load()
        
        # Model not loaded: return fallback
        if self._session is None:
            return image.convert("RGBA")
        
        # To numpy
        img = np.array(image)
        
        # Preprocess
        img_input, original_size = self._preprocess(img)
        
        try:
            pred = self._session.run([self._output_name], {self._input_name: img_input})[0]
        except Exception as e:
            # GPU failed, try CPU
            if hasattr(self, '_providers') and 'CUDAExecutionProvider' in self._providers:
                print(f"[RMBGModel] GPU inference failed (OOM), switching to CPU...")
                
                try:
                    # Release session
                    self._session = None
                    
                    # New CPU session
                    session_options = ort.SessionOptions()
                    session_options.log_severity_level = 3
                    
                    self._session = ort.InferenceSession(
                        self.model_path,
                        providers=['CPUExecutionProvider'],
                        sess_options=session_options
                    )
                    self._providers = ['CPUExecutionProvider']
                    
                    # Retry
                    pred = self._session.run([self._output_name], {self._input_name: img_input})[0]
                    print("[RMBGModel] CPU inference successful")
                    
                except Exception as e2:
                    print(f"[RMBGModel] CPU inference also failed: {e2}")
                    print("[RMBGModel] Falling back to no background removal")
                    return image.convert("RGBA")
            else:
                print(f"[RMBGModel] Inference failed: {e}, using fallback (no background removal)")
                return image.convert("RGBA")
        
        # Postprocess alpha
        alpha = self._postprocess(pred, original_size)
        
        # Merge alpha -> RGBA
        img_rgba = cv2.cvtColor(img, cv2.COLOR_RGB2RGBA)
        img_rgba[:, :, 3] = alpha
        
        # To PIL
        return Image.fromarray(img_rgba)