main.py

from PIL import Image
import os
import os.path
import numpy as np
from tqdm import tqdm

def process(input_dir, output_dir):
    # be careful with using this function, it will consume memory, access to disk and time
    images = []
    for f in tqdm(os.listdir(f"./{input_dir}")):
        try:
            with open(os.path.join(input_dir, f), "rb") as image:
                img_arr: bytes = image.read(h * w)
                while len(img_arr) > 0:
                    if len(img_arr) == h * w and img_arr not in images:
                        images.append(img_arr)
                    img_arr = image.read(h * w)
        except FileNotFoundError:
            continue

    rm_rf_dir(f"./{output_dir}")
    # And you can save them into png files
    count = 0
    for img in tqdm(images):
        png = Image.fromarray(np.reshape(list(img), (h, w)).astype("float32"), mode="L")
        png.save(f"./{output_dir}/{output_dir}_%d.png" % count)
        count += 1

def rm_rf_dir(dir):
    for file in os.listdir(dir):
        os.remove(os.path.join(dir, file))


h = 256  # height of image
w = 256  # width of image

if input("Should we do rats? (y/n) ") == "y":
    process("rats", "malware")
if input("Should we do not rats? (y/n) ") == "y":
    process("not_rats", "safe")