射影変換用の4点アノテーションツールをPythonで作る｜ベース画像の牌配置領域を指定する

import sys
import os
from pathlib import Path

from PyQt6.QtCore import Qt, QPointF
from PyQt6.QtGui import (
    QAction,
    QColor,
    QKeySequence,
    QPainter,
    QPen,
    QPixmap,
    QPolygonF,
)
from PyQt6.QtWidgets import (
    QApplication,
    QFileDialog,
    QGraphicsScene,
    QGraphicsView,
    QMainWindow,
)


# ============================================================
# Scene
# ============================================================

class AnnotationScene(QGraphicsScene):
    def __init__(self):
        super().__init__()
        self.image_item = None
        self.image_path = None
        self.txt_path = None
        self.current_points = []
        self.temp_point_items = []
        self.temp_line_items = []
        self.trapezoids = []


    def clear_annotation(self):
        self.clear()
        self.image_item = None
        self.current_points.clear()
        self.temp_point_items.clear()
        self.temp_line_items.clear()
        self.trapezoids.clear()


    def load_image(self, image_path):
        self.clear_annotation()
        self.image_path = image_path
        self.txt_path = str(Path(image_path).with_suffix(".txt"))
        pixmap = QPixmap(image_path)
        self.image_item = self.addPixmap(pixmap)
        self.setSceneRect(self.itemsBoundingRect())
        self.load_annotation()


    def load_annotation(self):
        if not self.txt_path:
            return
        if not os.path.exists(self.txt_path):
            return

        with open(self.txt_path, "r") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue

                vals = list(map(float, line.split(",")))
                if len(vals) != 8:
                    continue

                points = []
                for i in range(0, 8, 2):
                    points.append((vals[i], vals[i + 1]))

                polygon = QPolygonF(
                    [QPointF(x, y) for x, y in points]
                )

                poly_item = self.addPolygon(
                    polygon,
                    QPen(QColor(0, 255, 0), 2)
                )

                self.trapezoids.append({
                    "points": points,
                    "polygon": poly_item
                })


    def save_annotation(self):
        if not self.txt_path:
            return

        with open(self.txt_path, "w") as f:
            for data in self.trapezoids:
                row = []
                for x, y in data["points"]:
                    row.extend([
                        int(round(x)),
                        int(round(y))
                    ])
                f.write(",".join(map(str, row)))
                f.write("\n")
        print("saved:", self.txt_path)


    def add_temp_point(self, x, y):
        item = self.addEllipse(x - 4, y - 4, 8, 8, QPen(QColor(255, 0, 0)),)
        self.temp_point_items.append(item)


    def add_temp_line(self):
        if len(self.current_points) < 2:
            return
        x1, y1 = self.current_points[-2]
        x2, y2 = self.current_points[-1]
        line = self.addLine(
            x1,
            y1,
            x2,
            y2,
            QPen(QColor(255, 255, 0), 2)
        )
        self.temp_line_items.append(line)


    def finalize_quad(self):
        polygon = QPolygonF(
            [
                QPointF(x, y)
                for x, y in self.current_points
            ]
        )
        poly_item = self.addPolygon(
            polygon,
            QPen(QColor(0, 255, 0), 2)
        )
        self.trapezoids.append({
            "points": self.current_points.copy(),
            "polygon": poly_item
        })
        for item in self.temp_point_items:
            self.removeItem(item)

        for item in self.temp_line_items:
            self.removeItem(item)

        self.current_points.clear()
        self.temp_point_items.clear()
        self.temp_line_items.clear()


    def undo(self):
        if len(self.current_points) > 0:
            self.current_points.pop()
            if len(self.temp_point_items) > 0:
                item = self.temp_point_items.pop()
                self.removeItem(item)

            if len(self.temp_line_items) > 0:
                item = self.temp_line_items.pop()
                self.removeItem(item)
            return

        if len(self.trapezoids) > 0:
            last = self.trapezoids.pop()
            self.removeItem(last["polygon"])


    def delete_last(self):
        if len(self.trapezoids) == 0:
            return
        last = self.trapezoids.pop()
        self.removeItem(last["polygon"])


    def mousePressEvent(self, event):
        if event.button() != Qt.MouseButton.LeftButton:
            return super().mousePressEvent(event)

        pos = event.scenePos()
        x = pos.x()
        y = pos.y()
        self.current_points.append((x, y))
        self.add_temp_point(x, y)
        self.add_temp_line()
        if len(self.current_points) == 4:
            self.finalize_quad()

        super().mousePressEvent(event)



class AnnotationView(QGraphicsView):
    def __init__(self, scene):
        super().__init__(scene)
        self.zoom_factor = 1.15
        self.setRenderHint(
            QPainter.RenderHint.Antialiasing
        )
        self.setDragMode(
            QGraphicsView.DragMode.ScrollHandDrag
        )
        self.setTransformationAnchor(
            QGraphicsView.ViewportAnchor.AnchorUnderMouse
        )


    def wheelEvent(self, event):
        if event.angleDelta().y() > 0:
            self.scale(self.zoom_factor, self.zoom_factor)
        else:
            self.scale(1 / self.zoom_factor, 1 / self.zoom_factor)



class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        folder = QFileDialog.getExistingDirectory(
            self,
            "jpgフォルダ選択",
        )
        if not folder:
            sys.exit()
        self.image_files = sorted(
            [
                str(p)
                for p in Path(folder).glob("*.jpg")
            ]
        )
        if len(self.image_files) == 0:
            print("jpg not found")
            sys.exit()

        self.current_index = 0
        self.scene = AnnotationScene()
        self.view = AnnotationView(self.scene)
        self.setCentralWidget(self.view)
        self.resize(1600, 1000)
        self.load_image(0)
        self.create_actions()


    def create_actions(self):
        save_action = QAction(self)
        save_action.setShortcut(
            QKeySequence("Ctrl+S")
        )
        save_action.triggered.connect(
            self.scene.save_annotation
        )
        self.addAction(save_action)

        save_action_mac = QAction(self)
        save_action_mac.setShortcut(
            QKeySequence("Meta+S")
        )
        save_action_mac.triggered.connect(
            self.scene.save_annotation
        )
        self.addAction(save_action_mac)

        undo_action = QAction(self)
        undo_action.setShortcut(
            QKeySequence("Ctrl+Z")
        )
        undo_action.triggered.connect(
            self.scene.undo
        )
        self.addAction(undo_action)

        undo_action_mac = QAction(self)
        undo_action_mac.setShortcut(
            QKeySequence("Meta+Z")
        )
        undo_action_mac.triggered.connect(
            self.scene.undo
        )
        self.addAction(undo_action_mac)
        
        # 次画像
        next_action = QAction(self)
        next_action.setShortcut("N")
        next_action.triggered.connect(
            self.next_image
        )
        self.addAction(next_action)

        # 前画像
        prev_action = QAction(self)
        prev_action.setShortcut("P")
        prev_action.triggered.connect(
            self.prev_image
        )
        self.addAction(prev_action)
        
        next_action.setShortcut("Right")
        prev_action.setShortcut("Left")


    def load_image(self, index):
        self.current_index = index
        image_path = self.image_files[index]
        self.scene.load_image(image_path)
        self.setWindowTitle(
            f"[{index+1}/{len(self.image_files)}] "
            f"{os.path.basename(image_path)}"
        )


    def next_image(self):
        if self.current_index >= len(self.image_files) - 1:
            return
        self.scene.save_annotation()
        self.load_image(self.current_index + 1)


    def prev_image(self):
        if self.current_index <= 0:
            return
        self.scene.save_annotation()
        self.load_image(self.current_index - 1)


    def closeEvent(self, event):
        self.scene.save_annotation()
        event.accept()


def main():
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())


if __name__ == "__main__":
    main()