File size: 1,784 Bytes
1900386 |
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 |
import requests
import os
from PIL import Image
import numpy as np
from rembg import remove
class preprcessInput:
def __init__(self):
self.o_width = None
self.o_height = None
self.o_image = None
self.t_width = None
self.t_height = None
self.t_image = None
self.save_path = None
def remove_bg(self, file_path: str):
self.save_path = file_path[:-3]+'.png'
pic = Image.open(file_path)
self.o_width = np.asarray(pic).shape[1]
self.o_height = np.asarray(pic).shape[0]
try:
self.o_channels = np.asarray(pic).shape[2]
except Exception as e:
print("Single channel image and error", e)
os.remove(file_path)
self.o_image = remove(pic)
self.o_image.save(self.save_path)
os.remove(self.save_path)
return np.asarray(self.o_image)
def transform(self, width=768, height=1024):
newsize = (width, height)
self.t_height = height
self.t_width = width
pic = self.o_image
img = pic.resize(newsize)
self.t_image = img
background = Image.new("RGBA", newsize, (255, 255, 255, 255))
background.paste(img, mask=img.split()[3]) # 3 is the alpha channel
self.save_path = self.save_path[:-3] + '.jpg'
background.convert('RGB').save(self.save_path, 'JPEG')
return np.asarray(background.convert('RGB'))
# USAGE OF THE CLASS
# preprocess = preprcessInput()
# for images in os.listdir('F:\\AI\\StableVITON-master\\datasets\\test\\cloth'):
# print(images)
# if images[-3:] == 'jpg':
# op = preprocess.remove_bg('F:\\AI\\StableVITON-master\\datasets\\test\\cloth\\'+images)
# arr = preprocess.transform(768, 1024)
|