画像ファイルをいるものといらないものに仕分ける Tkinter 製 GUI ツール

image-classification-helper.py

#!/usr/bin/env python3

import os
import sys
import glob
import shutil
import tkinter as tk
from tkinter import messagebox
from tkinter import filedialog
from PIL import Image, ImageTk, ImageDraw

def alt_path(path, suffix="+"):
    while os.path.lexists(path):
        root, ext = os.path.splitext(path)
        head, tail = os.path.split(root)
        path = os.path.join(head, tail + suffix) + ext
    return path

def contain_size(img_width, img_height, frame_width, frame_height):
    if img_width / img_height < frame_width / frame_height:
        return max(1, round(img_width * (frame_height / img_height))), frame_height
    else:
        return frame_width, max(1, round(img_height * (frame_width / img_width)))

class App(tk.Tk):

    def __init__(self, src=None, dest=None):
        super().__init__(className="Image Classification Helper")
        self.title("Image Classification Helper")
        self.minsize(640, 480)
        self.geometry("720x540")
        self.grid_rowconfigure(1, weight=1)
        self.grid_columnconfigure(3, weight=1)
        self.icon = Image.new("RGBA", (1024, 1024))
        ImageDraw.Draw(self.icon).pieslice((48, 48, 986, 986), start=300, end=270, fill=(85, 192, 0), outline=(162, 192, 106))
        self.iconphoto(True, ImageTk.PhotoImage(self.icon.resize((128, 128), Image.LANCZOS)))

        while True:
            try:
                self.source_dir = src or filedialog.askdirectory(parent=self, title="Image Source Selection")
                src = None
            except:
                return
            if not self.source_dir:
                self.destroy()
                return
            exts = ["png", "jpg", "jpeg", "gif", "bmp", "tif", "tiff", "webp"]
            g = os.path.join(glob.escape(self.source_dir), "**", "*")
            self.images = sum([glob.glob(f"{g}.{e}", recursive=True) for e in exts], [])
            if not self.images:
                messagebox.showwarning("Empty Source", "No images found in the source directory!")
            else:
                self.images.reverse()
                break
        self.undo_funcs = []

        while True:
            try:
                dest_dir = dest or filedialog.askdirectory(parent=self, title="Destination Selection", mustexist=False)
                dest = None
            except:
                return
            if not dest_dir:
                self.destroy()
                return
            try:
                self.include_dir = os.path.normpath(os.path.join(dest_dir, "include"))
                self.exclude_dir = os.path.normpath(os.path.join(dest_dir, "exclude"))
                self.error_dir = os.path.normpath(os.path.join(dest_dir, "error"))
                os.makedirs(self.include_dir, exist_ok=True)
                os.makedirs(self.exclude_dir, exist_ok=True)
                os.makedirs(self.error_dir, exist_ok=True)
                break
            except:
                messagebox.showwarning("Error", "Can't create the destination directory!")

        self.include_button = tk.Button(self, text="Exclude (Left)", command=self.exclude)
        self.include_button.grid(row=0, column=0)
        self.exclude_button = tk.Button(self, text="Include (Right)", command=self.include)
        self.exclude_button.grid(row=0, column=1)
        self.skip_button = tk.Button(self, text="Skip (Control-s)", command=self.skip)
        self.skip_button.grid(row=0, column=2)
        self.undo_button = tk.Button(self, text="Undo (Control-z)", command=self.undo)
        self.undo_button.grid(row=0, column=4)

        self.cf = tk.Frame(bg="gray", bd=1)
        self.cf.grid(row=1, column=0, columnspan=5, sticky="nsew")
        self.canvas = tk.Canvas(self.cf, bg="green")
        self.canvas.pack(side="top", expand=True, fill="both")
        self.canvas.image = None
        self.canvas_image = self.canvas.create_image(0, 0, anchor="nw")

        self.progress = tk.Label(self, text="")
        self.progress.grid(row=2, column=0, columnspan=3, sticky="nw")
        self.file = tk.Label(self, text="")
        self.file.grid(row=2, column=3, columnspan=2, sticky="ne")

        self.bind("<Left>", self.exclude)
        self.bind("<Right>", self.include)
        self.bind("<Control-s>", self.skip)
        self.bind("<Control-z>", self.undo)
        self.bind("<Configure>", self.on_resize)
        self.next()

    def next(self):
        if not self.images:
            messagebox.showinfo("Finished", "No image left")
            self.destroy()
            return
        self.image = self.images.pop()
        try:
            self.canvas.image = Image.open(self.image).convert("RGB")
            self.update_image()
            left = len(self.images) + 1
            self.progress["text"] = f"{left} image{'' if left == 1 else 's'} left"
            self.file["text"] = self.image
        except:
            messagebox.showerror("Error", f"Can't open image file '{self.image}'")
            self.move_to_error()
            self.next()

    def undo(self, event=None):
        if len(self.undo_funcs) > 0:
            f = self.undo_funcs.pop()
        else:
            return
        try:
            image = f()
        except:
            messagebox.showerror("Error", "Rollback failed")
        else:
            self.images.append(self.image)
            self.images.append(image)
            self.next()

    def include(self, event=None):
        try:
            image = self.image
            filename = os.path.basename(image)
            filepath = alt_path(os.path.join(self.include_dir, filename))
            os.rename(image, filepath)
            self.undo_funcs.append(lambda: (os.rename(filepath, image), image)[1])
        except:
            messagebox.showerror("Error", f"Can't move image file '{image}'")
            self.move_to_error()
        self.next()

    def exclude(self, event=None):
        try:
            image = self.image
            filename = os.path.basename(image)
            filepath = alt_path(os.path.join(self.exclude_dir, filename))
            os.rename(image, filepath)
            self.undo_funcs.append(lambda: (os.rename(filepath, image), image)[1])
        except:
            messagebox.showerror("Error", f"Can't move image file '{image}'")
            self.move_to_error()
        self.next()

    def skip(self, event=None):
        image = self.image
        self.undo_funcs.append(lambda: image)
        self.next()

    def move_to_error(self):
        try:
            shutil.move(self.image, self.error_dir)
        except:
            messagebox.showerror("Error", f"Can't shunt problematic image file '{self.image}'")

    def update_image(self):
        if self.canvas.image is not None:
            width = self.canvas.winfo_width()
            height = self.canvas.winfo_height()
            w, h = contain_size(*self.canvas.image.size, width, height)
            self.canvas.photo = ImageTk.PhotoImage(self.canvas.image.resize((w, h), Image.BICUBIC))
            self.canvas.itemconfig(self.canvas_image, image=self.canvas.photo)

    def on_resize(self, event):
        self.update_image()

src = sys.argv[1] if len(sys.argv) > 1 else None
dest = sys.argv[2] if len(sys.argv) > 2 else None
App(src, dest).mainloop()