Python script for manually tagging images into binary classes

mark_n_move.py

#!/usr/bin/python

import argparse
import os
import shutil
from os.path import isfile, join

import cv2
import numpy as np


def info_print(idx):
    if idx < len(fpaths):
        print("Image", idx + 1, "of", len(fpaths), ":", fpaths[i]['fname'])
    else:
        print("That was the last of them! Well done! You deserve a cookie <3")


def move_image_to_folder(idx, mark):
    new_path = join(path, mark)
    print("Moving image", i, "to", new_path)
    if fpaths[idx]['path'] != new_path:
        shutil.move(join(fpaths[i]['path'], fpaths[i]['fname']), new_path)
        fpaths[idx]['path'] = new_path
    else:
        print("Same folder. Image wasn't moved")


if __name__ == "__main__":

    # Parse command line argument.
    parser = argparse.ArgumentParser()
    parser.add_argument("path", help="provide relative path to a folder containing images to classify.")
    args = parser.parse_args()
    path = join(os.getcwd(), args.path)

    # Get file paths.
    fpaths = [{"path": path, "fname": f} for f in os.listdir(path) if isfile(join(path, f))]
    print("no of fpaths: ", len(fpaths))

    folders = ['true', 'false']
    # Create folders if they don't exist.
    for folder in folders:
        if not os.path.exists(join(path, folder)):
            os.makedirs(join(path, folder))

    i = 0
    img = cv2.imread(join(fpaths[i]['path'], fpaths[i]['fname']))
    info_print(i)

    while i < len(fpaths):

        cv2.imshow('Classify image', img)
        img = cv2.imread(join(fpaths[i]['path'], fpaths[i]['fname']))

        k = cv2.waitKey(33)  # Detect key strokes.

        # Mark image as false.
        if k == ord('a'):
            move_image_to_folder(i, "false/")
            i += 1
            info_print(i)
            # Blink red to indicate the image was marked false.
            cv2.rectangle(img, tuple(np.subtract(img.shape[:2], (20, 20))), img.shape[:2], (0, 0, 255), -1)

        # Mark image as true.
        if k == ord('d'):
            move_image_to_folder(i, "true/")
            i += 1
            info_print(i)
            # Blink green to indicate the image was marked true.
            cv2.rectangle(img, tuple(np.subtract(img.shape[:2], (20, 20))), img.shape[:2], (0, 255, 0), -1)

        # Go back to previous image to correct the classification.
        if k == ord('s'):
            i -= 1
            info_print(i)

        # Esc key to stop.
        if k == 27:
            break