Manual classification to a target directory to "good" and "bad" images

manual-classification.py

""" UI to manually class bad/good images """
import os
import sys
from glob import glob

import cv2 as cv

directory = sys.argv[1]
if not directory:
    sys.exit(1)

dest = sys.argv[2]
if not dest:
    sys.exit(1)

bad = os.path.join(dest, "bad")
good = os.path.join(dest, "good")
files = glob(os.path.join(directory, "*"))
fc = len(files)

MAX_H = 900
MAX_W = 900
COUNT = 0

total = len(files)
while COUNT < total:
    f = files[COUNT]
    COUNT += 1

    # avoid good images
    bname = os.path.basename(f)
    is_good = os.path.join(dest, "good", bname)
    if os.path.exists(is_good):
        continue

    # open the image
    im = cv.imread(f)

    # rescale
    f1 = MAX_W / im.shape[1]
    f2 = MAX_H / im.shape[0]
    fac = min(f1, f2)  # resizing factor
    dim = (int(im.shape[1] * fac), int(im.shape[0] * fac))
    im = cv.resize(im, dim)

    # draw information on image
    gc = len(glob(os.path.join(good, "*")))
    bc = len(glob(os.path.join(bad, "*")))

    cv.putText(
        im,
        f"good: {gc}, bad: {bc} - {COUNT}/{fc}",
        (10, 50),
        cv.FONT_HERSHEY_SIMPLEX,
        0.75,
        (210, 224, 210),
        2,
    )

    # display and wait key
    cv.imshow("Press B for bad, G for good", im)
    k = cv.waitKey(0) & 0xFF
    if k == ord("b"):
        os.system(f"mv {f} {bad}")

    if k == ord("g"):
        os.system(f"mv {f} {good}")

    if k == ord("n"):
        continue

    if k == ord("d"):
        os.system(f"rm -f {f}")
        continue

    if k == ord("p"):
        COUNT -= 2
        continue

    if k == ord("q"):
        sys.exit(0)