70 lines
2.0 KiB
Python
70 lines
2.0 KiB
Python
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)
|
|
|
|
|
|
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)
|