diff --git a/.gitignore b/.gitignore index b6e4761..e265c19 100644 --- a/.gitignore +++ b/.gitignore @@ -127,3 +127,8 @@ dmypy.json # Pyre type checker .pyre/ + +# Custom +.idea/ +share/ +bin/ diff --git a/main.py b/main.py new file mode 100644 index 0000000..bd6a971 --- /dev/null +++ b/main.py @@ -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) diff --git a/pyvenv.cfg b/pyvenv.cfg new file mode 100644 index 0000000..4638ea8 --- /dev/null +++ b/pyvenv.cfg @@ -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