Simply inference logic

app.py

@@ -1,14 +1,17 @@
 import os
-from glob import glob
 import cv2
 import numpy as np
-from PIL import Image
 import torch
-from torchvision import transforms
-from transformers import AutoModelForImageSegmentation
 import gradio as gr
 import spaces
+
+from glob import glob
+from typing import Optional, Tuple
+
+from PIL import Image
 from gradio_imageslider import ImageSlider
+from transformers import AutoModelForImageSegmentation
+from torchvision import transforms
 
 torch.set_float32_matmul_precision('high')
 torch.jit.script = lambda f: f
@@ -16,21 +19,21 @@ torch.jit.script = lambda f: f
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
-def array_to_pil_image(image, size=(1024, 1024)):
+def array_to_pil_image(image: np.ndarray, size: Tuple[int, int] = (1024, 1024)) -> Image.Image:
     image = cv2.resize(image, size, interpolation=cv2.INTER_LINEAR)
     image = Image.fromarray(image).convert('RGB')
     return image
 
 
 class ImagePreprocessor():
-    def __init__(self, resolution=(1024, 1024)) -> None:
+    def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
         self.transform_image = transforms.Compose([
             # transforms.Resize(resolution),    # 1. keep consistent with the cv2.resize used in training 2. redundant with that in path_to_image()
             transforms.ToTensor(),
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
         ])
 
-    def proc(self, image):
+    def proc(self, image: Image.Image) -> torch.Tensor:
         image = self.transform_image(image)
         return image
 
@@ -45,14 +48,17 @@ usage_to_weights_file = {
     'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs'
 }
 
-from transformers import AutoModelForImageSegmentation
 birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True)
 birefnet.to(device)
 birefnet.eval()
 
 
 @spaces.GPU
-def predict(image, resolution, weights_file):
+def predict(
+    image: np.ndarray,
+    resolution: str,
+    weights_file: Optional[str]
+) -> Tuple[np.ndarray, np.ndarray]:
     global birefnet
     # Load BiRefNet with chosen weights
     _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
@@ -62,32 +68,32 @@ def predict(image, resolution, weights_file):
     birefnet.eval()
 
     resolution = f"{image.shape[1]}x{image.shape[0]}" if resolution == '' else resolution
-    # Image is a RGB numpy array.
     resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
-    images = [image]
-    image_shapes = [image.shape[:2] for image in images]
-    images = [array_to_pil_image(image, resolution) for image in images]
+    
+    image_shape = image.shape[:2]
+    image_pil = array_to_pil_image(image, tuple(resolution))
 
-    image_preprocessor = ImagePreprocessor(resolution=resolution)
-    images_proc = []
-    for image in images:
-        images_proc.append(image_preprocessor.proc(image))
-    images_proc = torch.cat([image_proc.unsqueeze(0) for image_proc in images_proc])
+    # Preprocess the image
+    image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
+    image_proc = image_preprocessor.proc(image_pil)
+    image_proc = image_proc.unsqueeze(0)
 
+    # Perform the prediction
     with torch.no_grad():
-        scaled_preds_tensor = birefnet(images_proc.to(device))[-1].sigmoid()   # BiRefNet needs an sigmoid activation outside the forward.
-    preds = []
-    for image_shape, pred_tensor in zip(image_shapes, scaled_preds_tensor):
-        if device == 'cuda':
-            pred_tensor = pred_tensor.cpu()
-        preds.append(torch.nn.functional.interpolate(pred_tensor.unsqueeze(0), size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy())
-    image_preds = []
-    for image, pred in zip(images, preds):
-        image = image.resize(pred.shape[::-1])
-        pred = np.repeat(np.expand_dims(pred, axis=-1), 3, axis=-1)
-        image_preds.append((pred * image).astype(np.uint8))
-
-    return image, image_preds[0]
+        scaled_pred_tensor = birefnet(image_proc.to(device))[-1].sigmoid()
+
+    if device == 'cuda':
+        scaled_pred_tensor = scaled_pred_tensor.cpu()
+    
+    # Resize the prediction to match the original image shape
+    pred = torch.nn.functional.interpolate(scaled_pred_tensor, size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy()
+
+    # Apply the prediction mask to the original image
+    image_pil = image_pil.resize(pred.shape[::-1])
+    pred = np.repeat(np.expand_dims(pred, axis=-1), 3, axis=-1)
+    image_pred = (pred * np.array(image_pil)).astype(np.uint8)
+
+    return image, image_pred
 
 
 examples = [[_] for _ in glob('examples/*')][:]