Files
test/examples/manual_annotator.py
2026-06-08 17:52:48 +08:00

658 lines
22 KiB
Python

"""
Manual annotation tool for camera calibration aids.
Use this when automatic corner detection is not practical. It lets you freeze a
camera/video frame or open an image, then manually mark corners, lines,
rectangles, and ROI boxes. Saved JSON coordinates are in original image pixels.
"""
import argparse
import json
import os
import sys
from datetime import datetime
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2
import numpy as np
def _fix_qt_plugin_path():
"""OpenCV can redirect Qt to cv2/qt/plugins; prefer PyQt5's plugins."""
try:
from PyQt5.QtCore import QLibraryInfo
os.environ["QT_QPA_PLATFORM_PLUGIN_PATH"] = QLibraryInfo.location(
QLibraryInfo.PluginsPath
)
except Exception:
path = os.environ.get("QT_QPA_PLATFORM_PLUGIN_PATH", "")
if "cv2" in path.replace("\\", "/"):
os.environ.pop("QT_QPA_PLATFORM_PLUGIN_PATH", None)
_fix_qt_plugin_path()
from PyQt5.QtCore import QPoint, QRect, Qt, QTimer
from PyQt5.QtGui import QColor, QFont, QImage, QPainter, QPen, QPixmap
from PyQt5.QtWidgets import (
QAction,
QApplication,
QButtonGroup,
QFileDialog,
QHBoxLayout,
QLabel,
QLineEdit,
QMainWindow,
QMessageBox,
QPushButton,
QSizePolicy,
QToolButton,
QVBoxLayout,
QWidget,
)
COLORS = {
"corner": QColor(0, 229, 255),
"line": QColor(255, 213, 79),
"rect": QColor(129, 199, 132),
"roi": QColor(255, 112, 67),
}
MODE_TEXT = {
"corner": "角点",
"line": "直线",
"rect": "矩形",
"roi": "ROI",
}
def cv_to_qpixmap(frame_bgr):
rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
h, w, ch = rgb.shape
qimg = QImage(rgb.data, w, h, ch * w, QImage.Format_RGB888).copy()
return QPixmap.fromImage(qimg)
def normalize_rect(p0, p1):
x0, y0 = p0
x1, y1 = p1
left, right = sorted((int(round(x0)), int(round(x1))))
top, bottom = sorted((int(round(y0)), int(round(y1))))
return [left, top, max(0, right - left), max(0, bottom - top)]
class ImageCanvas(QLabel):
def __init__(self, parent=None):
super().__init__(parent)
self.setMouseTracking(True)
self.setFocusPolicy(Qt.StrongFocus)
self.setAlignment(Qt.AlignCenter)
self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
self.setMinimumSize(640, 480)
self.frame = None
self.pixmap_src = None
self.annotations = []
self.mode = "corner"
self.drag_start = None
self.drag_current = None
self.hover_img_pos = None
self.status_callback = None
def set_status_callback(self, callback):
self.status_callback = callback
def set_frame(self, frame_bgr):
self.frame = frame_bgr.copy()
self.pixmap_src = cv_to_qpixmap(self.frame)
self.update()
def clear_annotations(self):
self.annotations.clear()
self.drag_start = None
self.drag_current = None
self.update()
def undo(self):
if self.annotations:
self.annotations.pop()
self.update()
def set_mode(self, mode):
self.mode = mode
self.drag_start = None
self.drag_current = None
self.update()
def image_size(self):
if self.frame is None:
return None
h, w = self.frame.shape[:2]
return [w, h]
def _target_rect(self):
if self.pixmap_src is None:
return QRect()
src_w = self.pixmap_src.width()
src_h = self.pixmap_src.height()
if src_w <= 0 or src_h <= 0:
return QRect()
widget_w = self.width()
widget_h = self.height()
scale = min(widget_w / src_w, widget_h / src_h)
draw_w = int(src_w * scale)
draw_h = int(src_h * scale)
x = (widget_w - draw_w) // 2
y = (widget_h - draw_h) // 2
return QRect(x, y, draw_w, draw_h)
def _widget_to_image(self, pos):
if self.pixmap_src is None:
return None
rect = self._target_rect()
if not rect.contains(pos):
return None
x = (pos.x() - rect.x()) * self.pixmap_src.width() / rect.width()
y = (pos.y() - rect.y()) * self.pixmap_src.height() / rect.height()
x = min(max(x, 0), self.pixmap_src.width() - 1)
y = min(max(y, 0), self.pixmap_src.height() - 1)
return [float(x), float(y)]
def _image_to_widget(self, point):
rect = self._target_rect()
x = rect.x() + point[0] * rect.width() / self.pixmap_src.width()
y = rect.y() + point[1] * rect.height() / self.pixmap_src.height()
return QPoint(int(round(x)), int(round(y)))
def mousePressEvent(self, event):
if event.button() != Qt.LeftButton or self.frame is None:
return
img_pos = self._widget_to_image(event.pos())
if img_pos is None:
return
if self.mode == "corner":
self.annotations.append(
{
"type": "corner",
"label": f"corner_{self._count_type('corner') + 1}",
"points": [[round(img_pos[0], 2), round(img_pos[1], 2)]],
}
)
self.update()
else:
self.drag_start = img_pos
self.drag_current = img_pos
self.update()
def mouseMoveEvent(self, event):
img_pos = self._widget_to_image(event.pos())
self.hover_img_pos = img_pos
if self.status_callback:
if img_pos is None:
self.status_callback(f"模式: {MODE_TEXT[self.mode]} | 坐标: -")
else:
self.status_callback(
f"模式: {MODE_TEXT[self.mode]} | 坐标: "
f"x={img_pos[0]:.1f}, y={img_pos[1]:.1f}"
)
if self.drag_start is not None and img_pos is not None:
self.drag_current = img_pos
self.update()
def mouseReleaseEvent(self, event):
if event.button() != Qt.LeftButton or self.drag_start is None:
return
img_pos = self._widget_to_image(event.pos()) or self.drag_current
if img_pos is None:
self.drag_start = None
self.drag_current = None
self.update()
return
p0 = [round(self.drag_start[0], 2), round(self.drag_start[1], 2)]
p1 = [round(img_pos[0], 2), round(img_pos[1], 2)]
if self.mode == "line":
self.annotations.append(
{
"type": "line",
"label": f"line_{self._count_type('line') + 1}",
"points": [p0, p1],
}
)
elif self.mode in ("rect", "roi"):
x, y, w, h = normalize_rect(p0, p1)
if w > 0 and h > 0:
self.annotations.append(
{
"type": self.mode,
"label": f"{self.mode}_{self._count_type(self.mode) + 1}",
"rect": [x, y, w, h],
"points": [[x, y], [x + w, y + h]],
}
)
self.drag_start = None
self.drag_current = None
self.update()
def keyPressEvent(self, event):
if event.key() == Qt.Key_Z and event.modifiers() & Qt.ControlModifier:
self.undo()
else:
super().keyPressEvent(event)
def paintEvent(self, event):
painter = QPainter(self)
painter.fillRect(self.rect(), QColor("#0b0b0b"))
if self.pixmap_src is None:
painter.setPen(QColor("#777777"))
painter.setFont(QFont("SimHei", 16))
painter.drawText(self.rect(), Qt.AlignCenter, "打开图片/视频或启动摄像头")
painter.end()
return
target = self._target_rect()
painter.drawPixmap(target, self.pixmap_src)
painter.setRenderHint(QPainter.Antialiasing, True)
for ann in self.annotations:
self._draw_annotation(painter, ann)
if self.drag_start is not None and self.drag_current is not None:
preview = self._preview_annotation()
if preview is not None:
self._draw_annotation(painter, preview, preview=True)
painter.end()
def _preview_annotation(self):
if self.mode == "line":
return {"type": "line", "points": [self.drag_start, self.drag_current]}
if self.mode in ("rect", "roi"):
x, y, w, h = normalize_rect(self.drag_start, self.drag_current)
return {
"type": self.mode,
"rect": [x, y, w, h],
"points": [[x, y], [x + w, y + h]],
}
return None
def _draw_annotation(self, painter, ann, preview=False):
color = COLORS.get(ann["type"], QColor(255, 255, 255))
pen = QPen(color, 2 if not preview else 1, Qt.DashLine if preview else Qt.SolidLine)
painter.setPen(pen)
painter.setBrush(Qt.NoBrush)
if ann["type"] == "corner":
p = self._image_to_widget(ann["points"][0])
painter.drawEllipse(p, 5, 5)
painter.drawLine(p.x() - 8, p.y(), p.x() + 8, p.y())
painter.drawLine(p.x(), p.y() - 8, p.x(), p.y() + 8)
self._draw_label(painter, p, ann.get("label", "corner"), color)
elif ann["type"] == "line":
p0 = self._image_to_widget(ann["points"][0])
p1 = self._image_to_widget(ann["points"][1])
painter.drawLine(p0, p1)
self._draw_label(painter, p0, ann.get("label", "line"), color)
elif ann["type"] in ("rect", "roi"):
x, y, w, h = ann["rect"]
p0 = self._image_to_widget([x, y])
p1 = self._image_to_widget([x + w, y + h])
rect = QRect(p0, p1).normalized()
painter.drawRect(rect)
self._draw_label(painter, rect.topLeft(), ann.get("label", ann["type"]), color)
def _draw_label(self, painter, pos, text, color):
painter.setFont(QFont("Consolas", 10))
metrics = painter.fontMetrics()
box = metrics.boundingRect(text).adjusted(-4, -2, 4, 2)
box.moveTopLeft(pos + QPoint(8, -box.height() - 4))
painter.fillRect(box, QColor(0, 0, 0, 170))
painter.setPen(color)
painter.drawText(box, Qt.AlignCenter, text)
def _count_type(self, ann_type):
return sum(1 for ann in self.annotations if ann["type"] == ann_type)
class MainWindow(QMainWindow):
def __init__(self, args):
super().__init__()
self.setWindowTitle("Orisys 手动标定标注工具")
self.args = args
self.capture = None
self.source = None
self.source_type = None
self.paused = True
self.timer = QTimer(self)
self.timer.timeout.connect(self._read_frame)
self._setup_ui()
self._setup_shortcuts()
if args.source:
self._open_source(args.source)
def _setup_ui(self):
root = QWidget()
self.setCentralWidget(root)
layout = QVBoxLayout(root)
layout.setContentsMargins(8, 8, 8, 8)
layout.setSpacing(8)
top = QHBoxLayout()
top.setSpacing(6)
self.source_edit = QLineEdit(str(self.args.source if self.args.source is not None else "0"))
self.source_edit.setPlaceholderText("摄像头编号、视频路径或图片路径")
top.addWidget(self.source_edit, 1)
self.btn_open = QPushButton("打开")
self.btn_open.clicked.connect(self._on_open_clicked)
top.addWidget(self.btn_open)
self.btn_file = QPushButton("选文件")
self.btn_file.clicked.connect(self._choose_file)
top.addWidget(self.btn_file)
self.btn_pause = QPushButton("暂停/继续")
self.btn_pause.clicked.connect(self._toggle_pause)
top.addWidget(self.btn_pause)
self.btn_save = QPushButton("保存 JSON")
self.btn_save.clicked.connect(self._save_json)
top.addWidget(self.btn_save)
self.btn_save_image = QPushButton("保存预览图")
self.btn_save_image.clicked.connect(self._save_preview)
top.addWidget(self.btn_save_image)
layout.addLayout(top)
tools = QHBoxLayout()
tools.setSpacing(6)
self.mode_group = QButtonGroup(self)
self.mode_group.setExclusive(True)
for mode in ("corner", "line", "rect", "roi"):
btn = QToolButton()
btn.setText(MODE_TEXT[mode])
btn.setCheckable(True)
btn.clicked.connect(lambda checked, m=mode: self.canvas.set_mode(m))
self.mode_group.addButton(btn)
tools.addWidget(btn)
if mode == "corner":
btn.setChecked(True)
tools.addStretch(1)
self.btn_undo = QPushButton("撤销")
self.btn_undo.clicked.connect(lambda: self.canvas.undo())
tools.addWidget(self.btn_undo)
self.btn_clear = QPushButton("清空")
self.btn_clear.clicked.connect(self._clear_annotations)
tools.addWidget(self.btn_clear)
layout.addLayout(tools)
self.canvas = ImageCanvas()
self.canvas.set_status_callback(self._set_status)
layout.addWidget(self.canvas, 1)
self.status = QLabel("模式: 角点 | 坐标: -")
self.status.setMinimumHeight(24)
layout.addWidget(self.status)
self.setStyleSheet(
"""
QMainWindow, QWidget { background: #111111; color: #dddddd; }
QLineEdit {
background: #0b0b0b; color: #dddddd; border: 1px solid #333333;
border-radius: 4px; padding: 5px 8px;
}
QPushButton, QToolButton {
background: #1d1d1d; color: #dddddd; border: 1px solid #333333;
border-radius: 4px; padding: 6px 12px;
}
QPushButton:hover, QToolButton:hover { background: #2a2a2a; }
QToolButton:checked {
background: #00323a; color: #00e5ff; border-color: #00a6b8;
}
QLabel { color: #aaaaaa; }
"""
)
def _setup_shortcuts(self):
shortcuts = [
("Ctrl+S", self._save_json),
("Ctrl+Z", self.canvas.undo),
("Space", self._toggle_pause),
("C", lambda: self._select_mode("corner")),
("L", lambda: self._select_mode("line")),
("B", lambda: self._select_mode("rect")),
("R", lambda: self._select_mode("roi")),
]
for key, callback in shortcuts:
action = QAction(self)
action.setShortcut(key)
action.triggered.connect(callback)
self.addAction(action)
def _select_mode(self, mode):
self.canvas.set_mode(mode)
for btn in self.mode_group.buttons():
if btn.text() == MODE_TEXT[mode]:
btn.setChecked(True)
break
def _set_status(self, text):
total = len(self.canvas.annotations)
self.status.setText(f"{text} | 标注数: {total}")
def _on_open_clicked(self):
self._open_source(self.source_edit.text().strip())
def _choose_file(self):
path, _ = QFileDialog.getOpenFileName(
self,
"选择图片或视频",
"data",
"媒体文件 (*.png *.jpg *.jpeg *.bmp *.tif *.tiff *.mp4 *.avi *.mov);;所有文件 (*)",
)
if path:
self.source_edit.setText(path)
self._open_source(path)
def _open_source(self, source):
self._release_capture()
self.canvas.clear_annotations()
self.source = source
src_for_cv = source
try:
src_for_cv = int(source)
except (TypeError, ValueError):
pass
if isinstance(src_for_cv, str) and self._is_image_file(src_for_cv):
frame = cv2.imread(src_for_cv, cv2.IMREAD_COLOR)
if frame is None:
self._show_error(f"无法读取图片: {src_for_cv}")
return
self.source_type = "image"
self.paused = True
self.canvas.set_frame(frame)
self._set_status("图片已打开")
return
cap = cv2.VideoCapture(src_for_cv)
if not cap.isOpened():
self._show_error(f"无法打开输入源: {source}")
return
self.capture = cap
self.source_type = "camera_or_video"
self.paused = False
self.timer.start(30)
self._read_frame()
def _read_frame(self):
if self.capture is None or self.paused:
return
ok, frame = self.capture.read()
if not ok:
self.paused = True
self.timer.stop()
self._set_status("视频结束或读取失败")
return
self.canvas.set_frame(frame)
def _toggle_pause(self):
if self.capture is None:
return
self.paused = not self.paused
if self.paused:
self.timer.stop()
else:
self.timer.start(30)
def _clear_annotations(self):
if not self.canvas.annotations:
return
reply = QMessageBox.question(
self,
"清空标注",
"确定清空当前所有标注?",
QMessageBox.Yes | QMessageBox.No,
QMessageBox.No,
)
if reply == QMessageBox.Yes:
self.canvas.clear_annotations()
def _save_json(self):
if self.canvas.frame is None:
self._show_error("当前没有图像可保存。")
return
default_dir = os.path.join("examples", "annotations")
os.makedirs(default_dir, exist_ok=True)
default_path = os.path.join(
default_dir, f"manual_annotation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
)
path, _ = QFileDialog.getSaveFileName(
self, "保存标注 JSON", default_path, "JSON (*.json);;所有文件 (*)"
)
if not path:
return
if not path.lower().endswith(".json"):
path += ".json"
payload = {
"created_at": datetime.now().isoformat(timespec="seconds"),
"source": self.source,
"source_type": self.source_type,
"image_size": self.canvas.image_size(),
"coordinate_system": "origin_top_left_pixels",
"annotations": self.canvas.annotations,
}
with open(path, "w", encoding="utf-8") as f:
json.dump(payload, f, ensure_ascii=False, indent=2)
self._set_status(f"已保存: {path}")
def _save_preview(self):
if self.canvas.frame is None:
self._show_error("当前没有图像可保存。")
return
preview = self._render_preview_image()
default_dir = os.path.join("examples", "annotations")
os.makedirs(default_dir, exist_ok=True)
default_path = os.path.join(
default_dir, f"manual_annotation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
)
path, _ = QFileDialog.getSaveFileName(
self, "保存标注预览图", default_path, "PNG (*.png);;所有文件 (*)"
)
if not path:
return
if not path.lower().endswith(".png"):
path += ".png"
cv2.imwrite(path, preview)
self._set_status(f"已保存预览图: {path}")
def _render_preview_image(self):
img = self.canvas.frame.copy()
for ann in self.canvas.annotations:
color = COLORS.get(ann["type"], QColor(255, 255, 255))
bgr = (color.blue(), color.green(), color.red())
if ann["type"] == "corner":
x, y = [int(round(v)) for v in ann["points"][0]]
cv2.drawMarker(img, (x, y), bgr, cv2.MARKER_CROSS, 18, 2)
cv2.circle(img, (x, y), 5, bgr, 2)
cv2.putText(img, ann.get("label", "corner"), (x + 8, y - 8),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, bgr, 1, cv2.LINE_AA)
elif ann["type"] == "line":
p0 = tuple(int(round(v)) for v in ann["points"][0])
p1 = tuple(int(round(v)) for v in ann["points"][1])
cv2.line(img, p0, p1, bgr, 2)
cv2.putText(img, ann.get("label", "line"), (p0[0] + 8, p0[1] - 8),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, bgr, 1, cv2.LINE_AA)
elif ann["type"] in ("rect", "roi"):
x, y, w, h = ann["rect"]
cv2.rectangle(img, (x, y), (x + w, y + h), bgr, 2)
cv2.putText(img, ann.get("label", ann["type"]), (x + 8, y + 18),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, bgr, 1, cv2.LINE_AA)
return img
def _release_capture(self):
self.timer.stop()
if self.capture is not None:
self.capture.release()
self.capture = None
def closeEvent(self, event):
self._release_capture()
event.accept()
@staticmethod
def _is_image_file(path):
return os.path.splitext(path.lower())[1] in {
".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"
}
def _show_error(self, message):
QMessageBox.warning(self, "错误", message)
self._set_status(message)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--source",
"-s",
default="0",
help="Camera index, video path, or image path. Default: 0",
)
return parser.parse_args()
def main():
args = parse_args()
QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)
QApplication.setAttribute(Qt.AA_UseHighDpiPixmaps, True)
app = QApplication(sys.argv)
window = MainWindow(args)
window.resize(1280, 860)
window.show()
sys.exit(app.exec_())
if __name__ == "__main__":
main()