Spaces:
Runtime error
Runtime error
File size: 4,905 Bytes
0dca4ae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 |
from PIL import Image,ImageOps
from rembg import new_session
import numpy as np
import cv2
class BackgroundRemover:
def __init__(self,model_name="unetp"):
self.session = new_session(model_name)
def read_image(self,path):
"""
Read the image.
Args:
path (str): The image path.
Returns:
PILImage: Original Image .
"""
image = Image.open(path)
return image
def resize_image(self,image,width,height):
"""
Resize the image.
Args:
img (PILImage): The image to be modified.
width (int): The width of the image.
height (int): The height of the image.
Returns:
PILImage: The resized image .
"""
image = image.resize((width,height))
return image
def putalpha_cutout(self,img, mask):
"""
Apply the specified mask to the image as an alpha cutout.
Args:
img (PILImage): The image to be modified.
mask (PILImage): The mask to be applied.
Returns:
PILImage: The modified image with the alpha cutout applied.
"""
img.putalpha(mask)
return img
def get_concat_v(self,img1, img2):
"""
Concatenate two images vertically.
Args:
img1 (PILImage): The first image.
img2 (PILImage): The second image to be concatenated below the first image.
Returns:
PILImage: The concatenated image.
"""
dst = Image.new("RGBA", (img1.width, img1.height + img2.height))
dst.paste(img1, (0, 0))
dst.paste(img2, (0, img1.height))
return dst
def get_concat_v_multi(self,imgs):
"""
Concatenate multiple images vertically.
Args:
imgs (List[PILImage]): The list of images to be concatenated.
Returns:
PILImage: The concatenated image.
"""
pivot = imgs.pop(0)
for im in imgs:
pivot = self.get_concat_v(pivot, im)
return pivot
def apply_background(self,img,filepath,w,h):
"""
Apply the specified background color to the image.
Args:
img (PILImage): The image to be modified.
color (Tuple[int, int, int, int]): The RGBA color to be applied.
Returns:
PILImage: The modified image with the background color applied.
"""
colored_image = self.read_image(filepath)
colored_image = self.resize_image(colored_image,width=w,height=h)
colored_image = colored_image.convert("RGBA")
colored_image.paste(img, mask=img)
return colored_image
def process_video(self,frame,background_path):
"""
Remove background to the image and replace with a selective background.
Args:
img_path: The first image path.
background_path: The background image path.
output_path :The directory where the file wile be saved. (optional)
save: Whether to save the file or not. (optional)
Returns:
PILImage: background removed image.
"""
lst_imgs=[]
input = Image.fromarray(frame)
width, height=input.size
input=ImageOps.exif_transpose(input)
masks = self.session.predict(input)
for mask in masks:
cutout = self.putalpha_cutout(input, mask)
lst_imgs.append(cutout)
cutout=input
if len(lst_imgs) > 0:
cutout = self.get_concat_v_multi(lst_imgs)
result = self.apply_background(input , background_path , width , height)
result = np.array(result)
return result
def process(self,img_path,background_path,output_path="",save=False):
"""
Remove background to the image and replace with a selective background.
Args:
img_path: The first image path.
background_path: The background image path.
output_path :The directory where the file wile be saved. (optional)
save: Whether to save the file or not. (optional)
Returns:
PILImage: background removed image.
"""
lst_imgs=[]
input = self.read_image(img_path)
width, height=input.size
input=ImageOps.exif_transpose(input)
masks = self.session.predict(input)
for mask in masks:
cutout = self.putalpha_cutout(input, mask)
lst_imgs.append(cutout)
cutout=input
if len(lst_imgs) > 0:
cutout = self.get_concat_v_multi(lst_imgs)
result = self.apply_background(input , background_path , width , height)
if save:
result.save(output_path)
return result
|