# from https://huggingface.co/spaces/eugenesiow/remove-bg/blob/main/app.py import cv2 import torch import numpy as np from torchvision import transforms class RemoveBackground(object): def __init__(self): self.model = torch.hub.load('pytorch/vision:v0.6.0', 'deeplabv3_resnet101', pretrained=True) self.model.eval() def make_transparent_foreground(self, pic, mask): # split the image into channels b, g, r = cv2.split(np.array(pic).astype('uint8')) # add an alpha channel with and fill all with transparent pixels (max 255) a = np.ones(mask.shape, dtype='uint8') * 255 # merge the alpha channel back alpha_im = cv2.merge([b, g, r, a], 4) # create a transparent background bg = np.zeros(alpha_im.shape) # setup the new mask new_mask = np.stack([mask, mask, mask, mask], axis=2) # copy only the foreground color pixels from the original image where mask is set foreground = np.where(new_mask, alpha_im, bg).astype(np.uint8) return foreground def remove_background(self, input_image): preprocess = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) input_tensor = preprocess(input_image) input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model # move the input and model to GPU for speed if available if torch.cuda.is_available(): input_batch = input_batch.to('cuda') self.model.to('cuda') with torch.no_grad(): output = self.model(input_batch)['out'][0] output_predictions = output.argmax(0) # create a binary (black and white) mask of the profile foreground mask = output_predictions.byte().cpu().numpy() background = np.zeros(mask.shape) bin_mask = np.where(mask, 255, background).astype(np.uint8) foreground = self.make_transparent_foreground(input_image, bin_mask) return foreground, bin_mask def inference(self, img): foreground, _ = self.remove_background(img) return foreground