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/*')][:]