task1.py

from functools import partial
import os
from multiprocesing import Process, Queue
from shutil import move
import time

import matplotlib.pyplot as plt
from pynput import mouse, Button
from imageio import imread


# global variables, so that we know them in `on_click`
# and don't need to pass everything with partial

# filename
CURRENT_FILE_NAME = None
# the 3 folders
IN_FOLDER = None
KEEP_FOLDER = None
DISCARD_FOLDER = None


# left click -> send to keep-folder
# right click -> send to discard-folder
def on_click(x, y, button, pressed, imqueue, ax):

    ax.clear()  # clear the currently shown image
    global CURRENT_FILE_NAME
    # move the current immage to keep or discard
    if button == Button.left:
        move(os.path.join(IN_FOLDER, CURRENT_FILE_NAME),
             os.path.join(KEEP_FOLDER, CURRENT_FILE_NAME))
    elif button == Button.right:
        move(os.path.join(IN_FOLDER, CURRENT_FILE_NAME),
             os.path.join(DISCARD_FOLDER, CURRENT_FILE_NAME))
    else:
        # do nothing if this button is unknown
        return None

    # stop the listener if the imqueue is empty and there are no more
    # files in the in-folder
    if imqueue.empty():
        if os.listdir(IN_FOLDER):  # we still have images -> workers were not fast enough
            return None
        else:  # no more images -> stop the listener
            return False
    # load the next image
    CURRENT_FILE_NAME, im = imqueue.get()

    # show the current image
    fig, ax = plt.subplots()
    ax.imshow(im)


def fill_fqueue(folder, fqueue):
    files = os.listdir(folder)
    for ff in files:
        fqueue.put(os.path.join(folder, ff))


# NOTE a worker will fail silently if the current file cannot be loaded with imread
def load_proc(fqueue, imqueue, max_queue_len, delta_load):
    while not fqueue.empty():
        # do we need this short sleep ?
        # reasoning: don't call qsize too often, don't know if this makes sense
        time.sleep(delta_load)
        if imqueue.qsize() >= max_queue_len:
            continue
        fname = fqueue.get()
        im = imread(fname)
        fname = os.path.split(fname)[1]
        imqueue.put((fname, im))


def main(in_folder, keep_folder, discard_folder,
         n_loaders=8, max_imqueue_len=100, delta_load=0.1):
    # make output folders
    os.makedirs(keep_folder, exist_ok=True)
    os.makedirs(discard_folder, exist_ok=True)

    # queue for filenames
    fqueue = Queue()
    fill_fqueue(in_folder, fqueue)

    # queue for the loaded images
    imqueue = Queue()

    # daemon processes to load images
    for _ in range(n_loaders):
        loader_p = Process(target=load_proc,
                           args=(fqueue, imqueue, max_imqueue_len, delta_load))
        loader_p.daemon = True
        loader_p.start()

    # assign global folder values
    global IN_FOLDER, KEEP_FOLDER, DISCARD_FOLDER
    IN_FOLDER = in_folder
    KEEP_FOLDER = keep_folder
    DISCARD_FOLDER = discard_folder

    fig, ax = plt.subplots()
    with mouse.Listener(on_click=partial(on_click,
                                         imqueue=imqueue,
                                         ax=ax)) as listener:
        listener.join()

        # show the first image
        global CURRENT_FILE_NAME
        CURRENT_FILE_NAME, im = imqueue.get()
        ax.imshow(im)
    fig.close()