60 lines
1.7 KiB
Python
60 lines
1.7 KiB
Python
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)
|