forked from SwayamInSync/clothes-virtual-try-on
-
Notifications
You must be signed in to change notification settings - Fork 0
/
remove_bg.py
62 lines (48 loc) · 1.75 KB
/
remove_bg.py
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
### Удаление фона и изменения размера
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('/content/inputs/test/image'):
print(images)
if images[-3:] == 'jpg':
op = preprocess.remove_bg(r'/content/inputs/test/image/'+images)
arr = preprocess.transform(768, 1024)