Pupil Center Labeling Tool

pupilcenter.py

from argparse import ArgumentParser
from pathlib import Path
from cv2 import waitKey
from cv2 import imshow
from cv2 import imread
from cv2 import namedWindow
from cv2 import setMouseCallback
from cv2 import destroyAllWindows
from cv2 import EVENT_LBUTTONUP

__name__ = 'Pupil Center Labeling Tool'

ix, iy = -1, -1
read_next = False

def ispressed(button, delay=1):
    return True if waitKey(delay) == button else False

def read_coordinates(event, x, y, flags, param):
    """
    Mouse callback function
    """
    global ix, iy
    global read_next
    if event == EVENT_LBUTTONUP:
        ix, iy = x, y
        read_next = True

parser = ArgumentParser(description='Tool for fast manual pupil center labeling.')

parser.add_argument('--dataset',
                    help='Path to dataset with eye images in PNG format.')
parser.add_argument('--storelabels',
                    help='Path to file for writing coordinates.',
                    default='./pupil_center_coordinates.txt')

args = parser.parse_args()
pictures = Path(args.dataset).glob('*.png')

data = 'Image_num X Y\n'

namedWindow(__name__)
setMouseCallback(__name__, read_coordinates)

picture_path = next(pictures)
picture_num = int(picture_path.name.split('.')[0])
image = imread(str(picture_path))
while not ispressed(27):
    if read_next:
        data += f'{picture_num} {ix} {iy}\n'
        read_next = False
        picture_path = next(pictures)
        picture_num = int(picture_path.name.split('.')[0])
        image = imread(str(picture_path))
    imshow(__name__, image)

with open(args.storelabels, 'w') as f:
    f.write(data)

destroyAllWindows()