Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Python script for manually tagging images into binary classes
#!/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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment