Created
March 6, 2019 23:20
-
-
Save constantinpape/1b76b1bd8215f6f98a20d22dddea77d4 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment