gather_images.py

desc = '''Script to gather data images with a particular label.

Usage: python gather_images.py <label_name> <num_samples>

The script will collect <num_samples> number of images and store them
in its own directory.

Only the portion of the image within the box displayed
will be captured and stored.

Press 'a' to start/pause the image collecting process.
Press 'q' to quit.

'''

import cv2
import os
import sys
import glob

try:
    label_name = sys.argv[1]
    num_samples = int(sys.argv[2])
except:
    print("Arguments missing.")
    print(desc)
    exit(-1)

IMG_SAVE_PATH = 'image_data'
IMG_CLASS_PATH = os.path.join(IMG_SAVE_PATH, label_name)
last_id = 0

try:
    os.mkdir(IMG_SAVE_PATH)
except FileExistsError:
    pass
try:
    os.mkdir(IMG_CLASS_PATH)
except FileExistsError:
    print("{} directory already exists.".format(IMG_CLASS_PATH))
    print("All images gathered will be saved along with existing items in this folder")
    list_of_files = glob.glob(os.path.join(IMG_CLASS_PATH, '*'))
    def comp(file):
        try :
            return int(os.path.splitext(file)[0].split('\\')[-1])
        except ValueError:
            return 0
    latest_file = max(list_of_files, key=comp)
    last_id = comp(latest_file)

cap = cv2.VideoCapture(0)

start = False
count = 0
reset_nb = 3
nb=0

while True:
    ret, frame = cap.read()
    if not ret:
        continue

    if count == num_samples:
        break

    cv2.rectangle(frame, (100, 100), (400, 400), (255, 255, 255), 2)

    if start:
        if nb==0 :
            roi = frame[100:400, 100:400]
            save_path = os.path.join(IMG_CLASS_PATH, '{}.jpg'.format(last_id + count + 1))
            cv2.imwrite(save_path, roi)
            count += 1
            nb = reset_nb
        else :
            nb -= 1
            

    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(frame, "Collecting {}".format(count),
            (50, 30), font, 0.7, (0, 0, 255), 2, cv2.LINE_AA)
    cv2.imshow("Collecting images", frame)

    k = cv2.waitKey(10)
    if k == ord('a'):
        start = not start

    if k == ord('q'):
        break

print("\n{} image(s) saved to {}".format(count, IMG_CLASS_PATH))
cap.release()
cv2.destroyAllWindows()