equalise_img_count.py

import os
import random

benDir = "../Equalised Images/Benign/"
malDir = "../Equalised Images/Malware/"

benFiles = os.listdir(benDir)
malFiles = os.listdir(malDir)

lb = len(benFiles)
lm = len(malFiles)

if lb > lm:
    rem = random.sample(benFiles, lb - lm)
    root = benDir
else:
    rem = random.sample(malFiles, lm - lb)
    root = malDir

for r in rem:
    os.remove(root + r)