crop_and_label.py

import cv2
import os
import glob
import argparse
from tqdm import tqdm
import pandas as pd


def click_and_crop(event, x, y, flags, param, points):
    if event == cv2.EVENT_LBUTTONUP:
        #points.append((x, y))
        points.append(x)
        points.append(y)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("directory")
    parser.add_argument("extension")
    args = parser.parse_args()

    print("q - quit")
    print("r - reset")
    print("u - undo")
    print("n - next")

    files = glob.glob("{}/*.{}".format(args.directory, args.extension))
    labels = pd.DataFrame(columns=["image", "points"])

    quit = False

    for f in tqdm(files):
        if quit:
            break

        image = cv2.imread(f)
        original = image.copy()
        points = []

        cv2.namedWindow("Image")
        cv2.setMouseCallback("Image", lambda event, x, y, flags, param: click_and_crop(event, x, y, flags, param, points))

        while True:
            if len(points) % 4 == 0:
                it = iter(points)
                for p in it:
                    cv2.rectangle(image, (p, next(it)), (next(it), next(it)), (255, 0, 0), 2)

            cv2.imshow("Image", image)
            key = cv2.waitKey(1) & 0xFF

            if key == ord("q"):
                quit = True
                break
            elif key == ord("r"):
                image = original.copy()
                points = []
            elif key == ord("u"):
                if len(points) > 0 and len(points) % 4 == 0:
                    image = original.copy()
                    points.pop()
                    points.pop()
                    points.pop()
                    points.pop()
            elif key == ord("n"):
                labels = labels.append({"image": os.path.basename(f), "points": " ".join([str(x) for x in points])}, ignore_index=True)
                labels.to_csv("labels.csv", index=False)
                break

        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()