YOLO用矩形アノテーションツールをPythonで自作する｜Tkinterでズーム・PANに対応

import copy
import re
import tkinter as tk
from dataclasses import dataclass
from pathlib import Path
from tkinter import filedialog, messagebox, simpledialog

from PIL import Image, ImageTk


CLASS_FILE = Path("./01_annotation/classes_yolo.txt")
DEFAULT_DIR = Path("./01_annotation/sample")

IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
COLORS = (
    "#e6194b", "#3cb44b", "#ffe119", "#4363d8",
    "#f58231", "#911eb4", "#46f0f0", "#f032e6",
    "#bcf60c", "#fabebe", "#008080", "#e6beff",
    "#333333",
)

MIN_BOX_SIZE = 5
MIN_ZOOM = 0.2
MAX_ZOOM = 10.0
ZOOM_STEP = 1.15


@dataclass
class Box:
    x1: float
    y1: float
    x2: float
    y2: float
    class_id: int

    def normalize(self) -> Box:
        return Box(
            min(self.x1, self.x2),
            min(self.y1, self.y2),
            max(self.x1, self.x2),
            max(self.y1, self.y2),
            self.class_id,
        )

    def clip(self, width: int, height: int) -> Box:
        box = self.normalize()
        return Box(
            max(0, min(box.x1, width)),
            max(0, min(box.y1, height)),
            max(0, min(box.x2, width)),
            max(0, min(box.y2, height)),
            box.class_id,
        )

    def contains(self, x: float, y: float) -> bool:
        box = self.normalize()
        return box.x1 <= x <= box.x2 and box.y1 <= y <= box.y2

    @property
    def width(self) -> float:
        return abs(self.x2 - self.x1)

    @property
    def height(self) -> float:
        return abs(self.y2 - self.y1)


class Annotator:
    def __init__(self, root: tk.Tk):
        self.root = root
        self.root.title("YOLO Annotator")
        self.root.geometry("1100x750")

        self.class_map = {}
        self.class_ids = []
        self.current_class = None

        self.image_paths = []
        self.index = 0
        self.image = None
        self.tk_image = None

        self.boxes = []
        self.history = []
        self.selected = None

        self.drawing = False
        self.draw_start = None
        self.draw_end = None

        self.zoom = 1.0
        self.pan_x = 0.0
        self.pan_y = 0.0
        self.pan_start = None
        self.pan_origin = None

        self.build_ui()
        self.load_classes()
        self.bind_events()

    def build_ui(self):
        self.canvas = tk.Canvas(self.root, bg="black", highlightthickness=0)
        self.canvas.pack(side="left", fill="both", expand=True)

        side = tk.Frame(self.root, width=220, padx=8, pady=8)
        side.pack(side="right", fill="y")
        side.pack_propagate(False)

        tk.Button(
            side,
            text="画像フォルダを開く",
            command=self.choose_folder,
        ).pack(fill="x")

        self.progress = tk.Label(side, text="画像未選択", justify="left")
        self.progress.pack(fill="x", pady=8)

        tk.Label(side, text="Classes", font=("", 12, "bold")).pack(anchor="w")

        self.class_list = tk.Listbox(side, exportselection=False)
        self.class_list.pack(fill="both", expand=True, pady=(4, 8))

        tk.Button(
            side,
            text="保存して次へ",
            command=self.save_and_next,
        ).pack(fill="x", pady=2)

        tk.Button(
            side,
            text="選択矩形を削除",
            command=self.delete_selected,
        ).pack(fill="x", pady=2)

        tk.Button(
            side,
            text="元に戻す",
            command=self.undo,
        ).pack(fill="x", pady=2)

        tk.Button(
            side,
            text="表示をリセット",
            command=self.reset_view,
        ).pack(fill="x", pady=2)

        self.status = tk.Label(
            side,
            text=(
                "左ドラッグ: 矩形作成\n"
                "右ドラッグ: PAN移動\n"
                "Ctrl + スクロール: ズーム"
            ),
            justify="left",
            wraplength=200,
        )
        self.status.pack(fill="x", pady=(10, 0))

    def bind_events(self):
        self.class_list.bind("<<ListboxSelect>>", self.select_class)

        self.canvas.bind("<ButtonPress-1>", self.left_press)
        self.canvas.bind("<B1-Motion>", self.left_drag)
        self.canvas.bind("<ButtonRelease-1>", self.left_release)

        # macOSでは右クリックがButton-2になる場合があるため両方に対応
        for button in (2, 3):
            self.canvas.bind(f"<ButtonPress-{button}>", self.pan_press)
            self.canvas.bind(f"<B{button}-Motion>", self.pan_drag)
            self.canvas.bind(f"<ButtonRelease-{button}>", self.pan_release)

        self.canvas.bind("<Control-MouseWheel>", self.zoom_image)
        self.canvas.bind("<Control-Button-4>", self.zoom_image)
        self.canvas.bind("<Control-Button-5>", self.zoom_image)
        self.canvas.bind("<Configure>", lambda _event: self.redraw())

        for key in ("<Control-s>", "<Command-s>"):
            self.root.bind(key, self.save_and_next)

        for key in ("<Control-z>", "<Command-z>"):
            self.root.bind(key, self.undo)

        self.root.bind("<Delete>", self.delete_selected)
        self.root.bind("<BackSpace>", self.delete_selected)

    def load_classes(self):
        if not CLASS_FILE.exists():
            messagebox.showerror(
                "エラー",
                f"クラスファイルが見つかりません。\n{CLASS_FILE.resolve()}",
            )
            return

        try:
            for line in CLASS_FILE.read_text(encoding="utf-8").splitlines():
                line = line.strip()
                if not line or line.startswith("#"):
                    continue

                class_id, name = line.split(maxsplit=1)
                class_id = int(class_id)

                self.class_map[class_id] = name
                self.class_ids.append(class_id)
                self.class_list.insert("end", f"{class_id}: {name}")
        except (OSError, ValueError) as error:
            messagebox.showerror("クラス読込エラー", str(error))
            return

        if self.class_ids:
            self.class_list.selection_set(0)
            self.current_class = self.class_ids[0]

    def select_class(self, _event=None):
        selection = self.class_list.curselection()
        if selection:
            self.current_class = self.class_ids[selection[0]]

    def choose_folder(self):
        initial_dir = DEFAULT_DIR if DEFAULT_DIR.exists() else Path.cwd()

        try:
            selected = filedialog.askdirectory(
                parent=self.root,
                title="画像フォルダを選択",
                initialdir=str(initial_dir),
                mustexist=True,
            )
        except tk.TclError:
            selected = simpledialog.askstring(
                "画像フォルダ",
                "画像フォルダの絶対パスを入力してください。",
                parent=self.root,
                initialvalue=str(initial_dir),
            )

        if not selected:
            return

        folder = Path(selected).expanduser()
        if not folder.is_dir():
            messagebox.showerror("エラー", f"フォルダが見つかりません。\n{folder}")
            return

        self.image_paths = sorted(
            [
                path
                for path in folder.iterdir()
                if path.suffix.lower() in IMAGE_EXTENSIONS
            ],
            key=lambda path: [
                int(part) if part.isdigit() else part.lower()
                for part in re.split(r"(\d+)", path.name)
            ],
        )

        if not self.image_paths:
            messagebox.showwarning("画像なし", "対応画像が見つかりません。")
            return

        self.index = 0
        self.load_image()

    def load_image(self):
        path = self.image_paths[self.index]

        try:
            with Image.open(path) as source:
                self.image = source.copy()
        except OSError as error:
            messagebox.showerror("画像読込エラー", str(error))
            return

        self.boxes = self.load_labels(path)
        self.history = [copy.deepcopy(self.boxes)]
        self.selected = None
        self.reset_view()

        self.progress.config(
            text=f"{self.index + 1} / {len(self.image_paths)}\n{path.name}"
        )

    def load_labels(self, image_path: Path):
        label_path = image_path.with_suffix(".txt")
        if not label_path.exists():
            return []

        width, height = self.image.size
        boxes = []

        try:
            for line in label_path.read_text(encoding="utf-8").splitlines():
                if not line.strip():
                    continue

                class_id, xc, yc, bw, bh = line.split()
                class_id = int(class_id)
                xc, yc, bw, bh = map(float, (xc, yc, bw, bh))

                boxes.append(
                    Box(
                        (xc - bw / 2) * width,
                        (yc - bh / 2) * height,
                        (xc + bw / 2) * width,
                        (yc + bh / 2) * height,
                        class_id,
                    ).clip(width, height)
                )
        except (OSError, ValueError) as error:
            messagebox.showwarning("ラベル読込エラー", str(error))

        return boxes

    def transform(self):
        if self.image is None:
            return 1.0, 0.0, 0.0

        canvas_w = max(self.canvas.winfo_width(), 1)
        canvas_h = max(self.canvas.winfo_height(), 1)
        image_w, image_h = self.image.size

        fit_scale = min(canvas_w / image_w, canvas_h / image_h)
        scale = fit_scale * self.zoom

        offset_x = (canvas_w - image_w * scale) / 2 + self.pan_x
        offset_y = (canvas_h - image_h * scale) / 2 + self.pan_y

        return scale, offset_x, offset_y

    def to_canvas(self, x, y):
        scale, offset_x, offset_y = self.transform()
        return x * scale + offset_x, y * scale + offset_y

    def to_image(self, x, y, clamp=False):
        if self.image is None:
            return None

        scale, offset_x, offset_y = self.transform()
        image_x = (x - offset_x) / scale
        image_y = (y - offset_y) / scale
        width, height = self.image.size

        if clamp:
            return (
                max(0, min(image_x, width)),
                max(0, min(image_y, height)),
            )

        if 0 <= image_x <= width and 0 <= image_y <= height:
            return image_x, image_y

        return None

    def redraw(self):
        self.canvas.delete("all")
        if self.image is None:
            return

        scale, offset_x, offset_y = self.transform()
        image_w, image_h = self.image.size
        display_w = max(1, round(image_w * scale))
        display_h = max(1, round(image_h * scale))

        resized = self.image.resize(
            (display_w, display_h),
            Image.Resampling.LANCZOS,
        )
        self.tk_image = ImageTk.PhotoImage(resized)

        self.canvas.create_image(
            offset_x,
            offset_y,
            anchor="nw",
            image=self.tk_image,
        )

        for index, box in enumerate(self.boxes):
            x1, y1 = self.to_canvas(box.x1, box.y1)
            x2, y2 = self.to_canvas(box.x2, box.y2)
            color = COLORS[box.class_id % len(COLORS)]

            self.canvas.create_rectangle(
                x1,
                y1,
                x2,
                y2,
                outline=color,
                width=4 if index == self.selected else 2,
            )
            self.canvas.create_text(
                x1,
                y1 - 4,
                text=f"{box.class_id}: {self.class_map.get(box.class_id, 'unknown')}",
                fill=color,
                anchor="sw",
            )

        if self.drawing and self.draw_start and self.draw_end:
            x1, y1 = self.to_canvas(*self.draw_start)
            x2, y2 = self.to_canvas(*self.draw_end)
            color = (
                COLORS[self.current_class % len(COLORS)]
                if self.current_class is not None
                else "white"
            )
            self.canvas.create_rectangle(
                x1,
                y1,
                x2,
                y2,
                outline=color,
                width=2,
                dash=(4, 2),
            )

    def left_press(self, event):
        point = self.to_image(event.x, event.y)
        if point is None:
            self.selected = None
            self.redraw()
            return

        self.selected = self.find_box(*point)
        if self.selected is not None:
            self.drawing = False
            self.redraw()
            return

        if self.current_class is None:
            self.set_status("クラスを選択してください")
            return

        self.drawing = True
        self.draw_start = point
        self.draw_end = point
        self.redraw()

    def left_drag(self, event):
        if not self.drawing:
            return

        self.draw_end = self.to_image(event.x, event.y, clamp=True)
        self.redraw()

    def left_release(self, event):
        if not self.drawing or self.draw_start is None:
            return

        self.drawing = False
        end = self.to_image(event.x, event.y, clamp=True)

        if end is None or self.current_class is None:
            self.draw_start = self.draw_end = None
            self.redraw()
            return

        width, height = self.image.size
        box = Box(
            *self.draw_start,
            *end,
            self.current_class,
        ).clip(width, height)

        self.draw_start = self.draw_end = None

        if box.width < MIN_BOX_SIZE or box.height < MIN_BOX_SIZE:
            self.set_status("小さすぎる矩形は追加しません")
            self.redraw()
            return

        self.boxes.append(box)
        self.selected = len(self.boxes) - 1
        self.history.append(copy.deepcopy(self.boxes))
        self.redraw()

    def find_box(self, x, y):
        for index in range(len(self.boxes) - 1, -1, -1):
            if self.boxes[index].contains(x, y):
                return index
        return None

    def pan_press(self, event):
        if self.image is None:
            return

        self.pan_start = (event.x, event.y)
        self.pan_origin = (self.pan_x, self.pan_y)

    def pan_drag(self, event):
        if self.pan_start is None or self.pan_origin is None:
            return

        self.pan_x = self.pan_origin[0] + event.x - self.pan_start[0]
        self.pan_y = self.pan_origin[1] + event.y - self.pan_start[1]
        self.redraw()

    def pan_release(self, _event):
        self.pan_start = None
        self.pan_origin = None

    def zoom_image(self, event):
        if self.image is None:
            return "break"

        image_point = self.to_image(event.x, event.y)
        if image_point is None:
            return "break"

        zoom_in = getattr(event, "num", None) == 4 or event.delta > 0
        factor = ZOOM_STEP if zoom_in else 1 / ZOOM_STEP
        new_zoom = max(MIN_ZOOM, min(self.zoom * factor, MAX_ZOOM))

        if new_zoom == self.zoom:
            return "break"

        self.zoom = new_zoom
        scale, _, _ = self.transform()

        canvas_w = max(self.canvas.winfo_width(), 1)
        canvas_h = max(self.canvas.winfo_height(), 1)
        image_w, image_h = self.image.size

        centered_x = (canvas_w - image_w * scale) / 2
        centered_y = (canvas_h - image_h * scale) / 2

        self.pan_x = event.x - image_point[0] * scale - centered_x
        self.pan_y = event.y - image_point[1] * scale - centered_y

        self.redraw()
        return "break"

    def reset_view(self):
        self.zoom = 1.0
        self.pan_x = 0.0
        self.pan_y = 0.0
        self.redraw()

    def delete_selected(self, _event=None):
        if self.selected is None:
            self.set_status("削除する矩形を選択してください")
            return "break"

        self.boxes.pop(self.selected)
        self.selected = None
        self.history.append(copy.deepcopy(self.boxes))
        self.redraw()
        return "break"

    def undo(self, _event=None):
        if len(self.history) <= 1:
            self.set_status("これ以上元に戻せません")
            return "break"

        self.history.pop()
        self.boxes = copy.deepcopy(self.history[-1])
        self.selected = None
        self.redraw()
        return "break"

    def save_and_next(self, _event=None):
        if self.image is None:
            self.set_status("画像フォルダを開いてください")
            return "break"

        width, height = self.image.size
        lines = []

        for box in self.boxes:
            box = box.clip(width, height)
            xc = ((box.x1 + box.x2) / 2) / width
            yc = ((box.y1 + box.y2) / 2) / height
            bw = box.width / width
            bh = box.height / height

            lines.append(
                f"{box.class_id} {xc:.6f} {yc:.6f} {bw:.6f} {bh:.6f}"
            )

        label_path = self.image_paths[self.index].with_suffix(".txt")

        try:
            label_path.write_text("\n".join(lines), encoding="utf-8")
        except OSError as error:
            messagebox.showerror("保存エラー", str(error))
            return "break"

        print(f"saved: {label_path}")

        if self.index + 1 >= len(self.image_paths):
            messagebox.showinfo("完了", "すべての画像を保存しました。")
            self.root.quit()
            return "break"

        self.index += 1
        self.load_image()
        return "break"

    def set_status(self, message):
        self.status.config(text=message)


if __name__ == "__main__":
    root = tk.Tk()
    Annotator(root)
    root.mainloop()