initial commit
This commit is contained in:
parent
3d9da24309
commit
dfc970393a
5
.gitignore
vendored
5
.gitignore
vendored
@ -127,3 +127,8 @@ dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# Custom
|
||||
.idea/
|
||||
share/
|
||||
bin/
|
||||
|
69
main.py
Normal file
69
main.py
Normal file
@ -0,0 +1,69 @@
|
||||
import os
|
||||
from PIL import Image, ImageChops
|
||||
from rembg import remove
|
||||
|
||||
|
||||
def scale_img(image, height):
|
||||
height_percent = (height / float(image.size[1]))
|
||||
width = int((float(image.size[0]) * float(height_percent)))
|
||||
return image.resize((width, height), Image.ANTIALIAS)
|
||||
|
||||
|
||||
input_path = './in/'
|
||||
convert_png = './.png/'
|
||||
removed_background = './.nobg/'
|
||||
output_full = './out/full/'
|
||||
output_scaled = './out/scaled/'
|
||||
image_suffix = ['png', 'jpg', 'jpeg', 'webp', 'gif']
|
||||
scales = [64, 128, 256]
|
||||
|
||||
|
||||
print('Converting input to png..')
|
||||
# Convert to png
|
||||
for i in os.listdir(input_path):
|
||||
if (i.split('.', -1)[1]) in image_suffix:
|
||||
img = Image.open(input_path + i)
|
||||
for s in image_suffix:
|
||||
if (i.split('.', -1))[1] == s:
|
||||
i = i.removesuffix(s)
|
||||
img.save(convert_png + i + '.png')
|
||||
|
||||
|
||||
print('Removing background with machine learning..')
|
||||
# Remove background
|
||||
for i in os.listdir(convert_png):
|
||||
if i.endswith('png'):
|
||||
img = Image.open(convert_png + i)
|
||||
nobg_img = remove(img)
|
||||
nobg_img.save(removed_background + i)q
|
||||
|
||||
|
||||
print('Cropping to content..')
|
||||
# Crop to content
|
||||
for i in os.listdir(removed_background):
|
||||
if i.endswith('png'):
|
||||
img = Image.open(removed_background + i)
|
||||
bg = Image.new(img.mode, img.size, img.getpixel((0, 0)))
|
||||
diff = ImageChops.difference(img, bg)
|
||||
diff = ImageChops.add(diff, diff, 0.2, -100)
|
||||
bbox = diff.getbbox()
|
||||
cropped = img.crop(bbox)
|
||||
cropped.save(output_full + i)
|
||||
|
||||
|
||||
print('Scaling..')
|
||||
# Scale
|
||||
for i in os.listdir(output_full):
|
||||
if i.endswith('png'):
|
||||
img = Image.open(output_full + i)
|
||||
for s in scales:
|
||||
img_scaled = scale_img(img, s)
|
||||
i = i.removesuffix('.png')
|
||||
img_scaled.save(output_scaled + i + '.' + str(s) + '.png')
|
||||
|
||||
# cleanup
|
||||
for f in [f for f in os.listdir(removed_background) if os.path.isfile(f)]:
|
||||
os.remove(removed_background + f)
|
||||
|
||||
for f in [f for f in os.listdir(convert_png) if os.path.isfile(f)]:
|
||||
os.remove(convert_png + f)
|
8
pyvenv.cfg
Normal file
8
pyvenv.cfg
Normal file
@ -0,0 +1,8 @@
|
||||
home = /usr
|
||||
implementation = CPython
|
||||
version_info = 3.9.9.final.0
|
||||
virtualenv = 20.13.0
|
||||
include-system-site-packages = false
|
||||
base-prefix = /usr
|
||||
base-exec-prefix = /usr
|
||||
base-executable = /usr/bin/python3.9
|
Loading…
x
Reference in New Issue
Block a user