Image Labeler

gui_labeler.py

import argparse
import csv
import os
import shutil
import sys
import tkinter as tk
from typing import TextIO, Tuple
from PIL import Image, ImageTk


class Labeler:
    def __init__(self, input_folder: str, output_folder: str, csvfile: TextIO, labels: int, resolution: Tuple[int, int]):
        self.root = tk.Tk()
        self.root.title('Labeler')
        self.root.configure(bg='black', padx=10)

        self.input_folder = input_folder
        self.output_folder = output_folder
        self.csvfile = csvfile
        self.writer = csv.writer(self.csvfile)
        self.labels = labels
        self.files = os.listdir(input_folder)
        self.resolution = resolution

        # buttons
        self.button_frame = tk.Frame(self.root)
        self.button_frame.pack(side='top', anchor='w')
        self.buttons = []
        for i in range(labels):
            button = tk.Button(
                self.button_frame,
                text=i,
                command=lambda j = i: self.move(j),
                bg='black',
                fg='white',
                padx=10,
                bd=4,
            )
            button.config(font=('Calibri', 12))
            button.pack(side='left')
            self.buttons.append(button)

        self.show_button = tk.Button(
            self.button_frame,
            text='Show',
            command=self.show,
            bg='black',
            fg='white',
            padx=10,
            bd=4,
        )
        self.show_button.pack(side='left')

        self.skip_button = tk.Button(
            self.button_frame,
            text='Skip',
            command=self.next,
            bg='black',
            fg='white',
            padx=10,
            bd=4,
        )
        self.skip_button.pack(side='left')

        # keyboard
        if self.labels <= 10:
            for i in range(self.labels):
                self.root.bind(f'<Key-{i}>', lambda e, j = i: self.move(j))
                self.root.bind(f'<Key-KP_{i}>', lambda e, j = i: self.move(j))
        
        self.index = 0
        self.filename = self.files[self.index]
        image = Image.open(os.path.join(self.input_folder, self.filename))
        width, height = image.size
        width_ratio = self.resolution[0] / width
        height_ratio = self.resolution[1] / height
        ratio = min(width_ratio, height_ratio)
        self.image = image.resize((int(width * ratio), int(height * ratio)), Image.LANCZOS)

        # image
        self.tk_image = ImageTk.PhotoImage(self.image)
        self.panel = tk.Label(self.root, image=self.tk_image, bg='black')
        self.panel.pack(side='bottom', fill='both', expand=1)

        # name
        self.label = tk.Label(self.root, text=self.filename, bg='black', fg='white')
        self.label.pack(side='bottom', fill='both', expand=1)

        # count
        self.remaining = tk.Label(self.root, text=f'{len(self.files) - self.index} left', bg='black', fg='white')
        self.remaining.pack(side='bottom', fill='both', expand=1)

    def show(self):
        print('skip')
        self.image.show()

    def move(self, label: int):
        print(f'{self.filename}: {label}')
        self.writer.writerow([self.filename, label])
        input_path = os.path.join(self.input_folder, self.filename)
        output_path = os.path.join(self.output_folder, str(label), self.filename)
        shutil.move(input_path, output_path)
        self.next()

    def next(self):
        self.index += 1
        if self.index >= len(self.files):
            self.panel.destroy()
            self.label.configure(text='FINISHED')
            self.remaining.configure(text='0 left')
        else:
            self.filename = self.files[self.index]
            image = Image.open(os.path.join(self.input_folder, self.filename))
            width, height = image.size
            width_ratio = self.resolution[0] / width
            height_ratio = self.resolution[1] / height
            ratio = min(width_ratio, height_ratio)
            self.image = image.resize((int(width * ratio), int(height * ratio)), Image.LANCZOS)
            self.tk_image = ImageTk.PhotoImage(self.image)
            self.panel.configure(image=self.tk_image)
            self.label.configure(text=self.filename)
            self.remaining.configure(text=f'{len(self.files) - self.index} left')

    def start(self):
        self.root.mainloop()


def get_args():
    parser = argparse.ArgumentParser(description='labeler')
    parser.add_argument('--input', type=str, required=True, help='input folder')
    parser.add_argument('--output', type=str, required=True, help='output folder')
    parser.add_argument('--csv', type=str, required=True, help='csv of labels')
    parser.add_argument('--labels', type=int, required=True, help='number of labels')
    parser.add_argument('--resolution', type=int, nargs=2, default=[1366, 768], help='resolution size')
    return parser.parse_args()


def main(args):
    if not os.path.isdir(args.input):
        sys.exit('input folder does not exist')

    # setup directories
    for i in range(args.labels):
        path = os.path.join(args.output, str(i))
        if not os.path.isdir(path):
            os.makedirs(path)

    # create csv file if does not exist
    if not os.path.isfile(args.csv):
        print('creating csv file')
        with open(args.csv, 'w') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(['filename', 'label'])

    with open(args.csv, 'a') as csvfile:
        gui = Labeler(args.input, args.output, csvfile, args.labels, args.resolution)
        gui.start()


if __name__ == '__main__':
    args = get_args()
    main(args)