pywai/main.py

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)