python
复制代码
import glob
import sys
import json
import os
import shutil
from PyQt5.QtWidgets import (
QApplication, QMainWindow, QTableWidget, QTableWidgetItem,
QSplitter, QVBoxLayout, QWidget, QPushButton, QRadioButton,
QButtonGroup, QLabel, QHBoxLayout, QMessageBox, QScrollArea
)
from PyQt5.QtCore import Qt, QUrl
from PyQt5.QtGui import QColor, QPixmap, QImage, QPainter, QPen, QKeyEvent
from natsort import natsorted
class CustomTableWidget(QTableWidget):
def __init__(self, parent=None):
super().__init__(parent)
self.parent_window = parent
def keyPressEvent(self, event):
"""重写表格的键盘事件处理"""
if event.key() == Qt.Key_Up or event.key() == Qt.Key_Down:
# 获取当前选中行
current_row = self.currentRow()
if event.key() == Qt.Key_Up and current_row > 0:
# 上箭头键，选择上一行
new_row = current_row - 1
self.selectRow(new_row)
self.setCurrentCell(new_row, 0)
if self.parent_window:
self.parent_window.display_image(new_row, 0)
elif event.key() == Qt.Key_Down and current_row < self.rowCount() - 1:
# 下箭头键，选择下一行
new_row = current_row + 1
self.selectRow(new_row)
self.setCurrentCell(new_row, 0)
if self.parent_window:
self.parent_window.display_image(new_row, 0)
# 数字键（主键盘和小键盘）监听
elif Qt.Key_0 <= event.key() <= Qt.Key_9:
num = event.key() - Qt.Key_0
if self.parent_window:
self.parent_window.handle_number_key(num)
elif Qt.KeypadModifier & event.modifiers() and Qt.Key_0 <= event.key() <= Qt.Key_9:
num = event.key() - Qt.Key_0
print(f"小键盘数字键按下: {num}")
if self.parent_window:
self.parent_window.handle_number_key(num)
else:
# 其他按键按默认方式处理
super().keyPressEvent(event)
class ImageAnnotator(QMainWindow):
def __init__(self, base_dir):
super().__init__()
self.setWindowTitle(f"图片标注可视化工具 - {os.path.basename(base_dir)}")
self.setGeometry(100, 100, 1400, 900)
# 启用拖放功能
self.setAcceptDrops(True)
# 主布局
main_widget = QWidget()
self.setCentralWidget(main_widget)
layout = QVBoxLayout()
main_widget.setLayout(layout)
# 添加当前目录显示
self.dir_label = QLabel(f"当前目录: {base_dir}")
self.dir_label.setStyleSheet("background-color: #e0e0e0; padding: 5px;")
self.dir_label.setFixedHeight(40) # 设置高度为40像素
self.dir_label.setTextInteractionFlags(Qt.TextSelectableByMouse) #启用鼠标选中文本复制
layout.addWidget(self.dir_label)
# 分割左右区域
splitter = QSplitter(Qt.Horizontal)
# 左侧：图片文件列表（表格）- 使用自定义表格控件
self.table = CustomTableWidget(self)
self.table.setColumnCount(3)
self.table.setHorizontalHeaderLabels(["图片文件", "状态", "标注数量"])
self.table.setEditTriggers(QTableWidget.NoEditTriggers)
self.table.cellClicked.connect(self.display_image)
self.table.setColumnWidth(0, 250)
self.table.setColumnWidth(1, 100)
self.table.setColumnWidth(2, 100)
# 设置表格可以获取焦点，以便接收键盘事件
self.table.setFocusPolicy(Qt.StrongFocus)
self.table.setSelectionBehavior(QTableWidget.SelectRows)
self.table.setSelectionMode(QTableWidget.SingleSelection)
# 右侧：图片显示和标注区域
right_panel = QWidget()
right_layout = QVBoxLayout()
# 图片显示区域
self.image_label = QLabel()
self.image_label.setAlignment(Qt.AlignCenter)
self.image_label.setMinimumSize(500, 400)
self.image_label.setStyleSheet("border: 1px solid gray; background-color: #f0f0f0;")
self.image_label.setText("请选择图片")
# 添加拖拽提示
self.image_label.setAcceptDrops(True)
# 创建滚动区域用于显示大图
scroll_area = QScrollArea()
scroll_area.setWidget(self.image_label)
scroll_area.setWidgetResizable(True)
self.pic_label=QLabel("图片预览（带标注框）：")
right_layout.addWidget(self.pic_label)
# right_layout.addWidget(scroll_area)
right_layout.addWidget(scroll_area, stretch=2)
# 标注选项
self.label_map = {0: "ok", 1: "err", 2: "other"}
# 创建分类目录
self.output_dirs = {}
self.base_dir = base_dir
self.create_output_dirs()
# 单选按钮组
self.radio_group = QButtonGroup()
right_layout.addWidget(QLabel("分类选项："))
# 创建单选按钮的垂直布局
radio_layout = QVBoxLayout()
radio_widget = QWidget()
radio_widget.setLayout(radio_layout)
for key, text in self.label_map.items():
radio = QRadioButton(f"{text} ({key})")
radio_layout.addWidget(radio)
self.radio_group.addButton(radio, key)
self.radio_group.buttonClicked[int].connect(self.on_radio_selected)
right_layout.addWidget(radio_widget)
# 显示目标目录信息
dir_info = QLabel("分类后图片和JSON将自动移动到对应目录：\n" +
"\n".join([f"{label_id}_{label_name}" for label_id, label_name in self.label_map.items()]))
dir_info.setStyleSheet("color: blue; font-size: 10px;")
right_layout.addWidget(dir_info)
# 添加拖拽提示
drag_info = QLabel("提示: 拖拽目录到此窗口可切换数据源")
drag_info.setStyleSheet("color: green; font-size: 10px; background-color: #f0f0f0; padding: 5px;")
right_layout.addWidget(drag_info)
# right_layout.addStretch()
right_layout.addStretch(1)
right_panel.setLayout(right_layout)
# 添加到布局
splitter.addWidget(self.table)
splitter.addWidget(right_panel)
splitter.setSizes([400, 800])#两列的宽度比例
layout.addWidget(splitter)
# 存储所有图片的原始信息
self.all_images = [] # 存储 (原始图片路径, 原始JSON路径, 新图片路径, 新JSON路径, 状态, 标注数量) 的元组
# 初始化图片文件列表
self.refresh_image_files()
self.load_image_files()
# 默认选中第一行
if self.all_images:
self.table.selectRow(0)
self.table.setFocus()
self.display_image(0, 0)
def create_output_dirs(self):
"""创建分类目录"""
self.output_dirs = {}
for label_id, label_name in self.label_map.items():
dir_path = os.path.join(self.base_dir, f"{label_id}_{label_name}")
# os.makedirs(dir_path, exist_ok=True)
self.output_dirs[label_id] = dir_path
def dragEnterEvent(self, event):
"""处理拖拽进入事件"""
if event.mimeData().hasUrls():
event.acceptProposedAction()
def dropEvent(self, event):
"""处理拖拽释放事件"""
urls = event.mimeData().urls()
if urls:
# 获取拖拽的第一个路径
path = urls[0].toLocalFile()
if os.path.isfile(path):
self.base_dir=os.path.dirname(path)
elif os.path.isdir(path):
# 更新基础目录
self.base_dir = path
else:
QMessageBox.warning(self, "错误", "请拖拽目录而非文件！")
return
# 更新窗口标题和目录标签
self.setWindowTitle(f"图片标注可视化工具 - {os.path.basename(self.base_dir)}")
self.dir_label.setText(f"当前目录: {self.base_dir}")
# 重新创建分类目录
self.create_output_dirs()
# 刷新图片列表
self.refresh_image_files()
self.load_image_files()
# 默认选中第一行
if self.all_images:
self.table.selectRow(0)
self.table.setFocus()
self.display_image(0, 0)
# 显示成功消息
self.statusBar().showMessage(f"已切换到目录: {self.base_dir}", 3000)
def refresh_image_files(self):
"""刷新图片文件列表，包含所有图片（包括已分类的）"""
self.all_images = []
img_files = ['%s/%s' % (i[0].replace("\\", "/"), j) for i in os.walk(self.base_dir) for j in i[-1] if
j.lower().endswith(('png', 'jpg', 'jpeg'))]
img_files = natsorted(img_files)
for img_path in img_files:
dir_name =os.path.basename(os.path.dirname(img_path))
json_path = self.find_json_file(img_path)
annotation_count = self.get_annotation_count(json_path)
if any(f"{label_id}_{self.label_map[label_id]}" in dir_name for label_id in self.label_map):
self.all_images.append \
((img_path, json_path, img_path, json_path, f"已分类: {dir_name}", annotation_count))
else:
self.all_images.append((img_path, json_path, None, None, "待分类", annotation_count))
def find_json_file(self, img_path):
"""查找与图片对应的JSON文件"""
# 尝试多种可能的JSON文件名
base_name = os.path.splitext(img_path)[0]
possible_json_paths = [
base_name + '.json',
base_name + '_hand.json',
base_name.replace('_s', '') + '.json', # 处理_s后缀的图片
]
for json_path in possible_json_paths:
if os.path.exists(json_path):
return json_path
return None
def get_annotation_count(self, json_path):
"""从JSON文件中获取标注数量"""
if not json_path or not os.path.exists(json_path):
return 0
try:
with open(json_path, 'r', encoding='utf-8') as f:
data = json.load(f)
return len(data.get('shapes', []))
except:
return 0
def move_image_and_json(self, img_path, json_path, label_id):
"""将图片和JSON文件移动到对应的分类目录"""
try:
if label_id not in self.output_dirs:
return None, None
target_dir = self.output_dirs[label_id]
os.makedirs(target_dir,exist_ok=True)
img_filename = os.path.basename(img_path)
# 移动图片文件
img_target_path = os.path.join(target_dir, img_filename)
counter = 1
base_name, ext = os.path.splitext(img_filename)
if os.path.exists(img_target_path):
json_target_path = None
if json_path and os.path.exists(json_path):
json_filename = os.path.basename(json_path)
json_target_path = os.path.join(target_dir, json_filename)
return img_target_path, json_target_path
shutil.move(img_path, img_target_path)
# 移动JSON文件（如果存在）
json_target_path = None
if json_path and os.path.exists(json_path):
json_filename = os.path.basename(json_path)
json_target_path = os.path.join(target_dir, json_filename)
shutil.move(json_path, json_target_path)
return img_target_path, json_target_path
except Exception as e:
print(f"移动文件失败: {e}")
return None, None
def on_radio_selected(self, checked_id):
"""单选按钮选择事件"""
current_row = self.table.currentRow()
if current_row >= 0 and current_row < len(self.all_images):
original_img, original_json, current_img, current_json, status, annotation_count = self.all_images[
current_row]
# checked_id = self.radio_group.checkedId()
if checked_id <0:
print('checked_id <0',checked_id)
return
# 使用当前路径（如果已移动）或原始路径
img_path = current_img if current_img else original_img
json_path = current_json if current_json else original_json
# 检查图片文件是否存在
if not os.path.exists(img_path):
QMessageBox.warning(self, "警告", "图片文件不存在！")
self.refresh_image_files()
self.load_image_files()
return
# 移动图片和JSON到对应目录
new_img_path, new_json_path = self.move_image_and_json(img_path, json_path, checked_id)
if new_img_path:
# 更新内存中的数据
label_name = self.label_map.get(checked_id, "未知")
label_show=f"{checked_id}_{label_name}"
self.all_images[current_row] = \
(
original_img, original_json, new_img_path, new_json_path, f"已分类: {label_show}", annotation_count)
# 更新表格显示
self.update_table_row(current_row)
# 显示成功消息
self.statusBar().showMessage(f"已分类: {os.path.basename(original_img)} -> {label_show}", 3000)
# 自动选择下一行（如果还有未分类的图片）
next_row = current_row+1
next_row = min(len(self.all_images)-1,max(0,next_row))
self.table.selectRow(next_row)
self.table.setCurrentCell(next_row, 0)
self.table.setFocus()
self.display_image(next_row, 0)
else:
QMessageBox.warning(self, "错误", "移动文件失败！")
def find_next_unclassified(self, start_row):
"""从指定行开始查找下一个未分类的图片"""
for i in range(start_row + 1, len(self.all_images)):
original_img, original_json, current_img, current_json, status, annotation_count = self.all_images[i]
if status == "待分类":
return i
# 如果后面没有未分类的，从开头找
for i in range(0, start_row):
original_img, original_json, current_img, current_json, status, annotation_count = self.all_images[i]
if status == "待分类":
return i
return -1 # 没有未分类的图片了
def update_table_row(self, row):
"""更新表格中指定行的显示状态"""
if 0 <= row < len(self.all_images):
original_img, original_json, current_img, current_json, status, annotation_count = self.all_images[row]
# 文件名列
file_item = QTableWidgetItem(os.path.basename(original_img))
# 状态列
status_item = QTableWidgetItem(status)
# 标注数量列
count_item = QTableWidgetItem(str(annotation_count))
# 设置背景色
if status == "待分类":
color = QColor(255, 200, 200) # 浅红色
elif "err" in status:
# color = QColor(255, 165, 0) # 浅橙色
color = QColor(255, 102, 102) # 浅橙色
else:
color = QColor(200, 255, 200) # 浅绿色
file_item.setBackground(color)
status_item.setBackground(color)
count_item.setBackground(color)
self.table.setItem(row, 0, file_item)
self.table.setItem(row, 1, status_item)
self.table.setItem(row, 2, count_item)
def load_image_files(self):
"""加载图片文件到表格"""
self.table.setRowCount(len(self.all_images))
for i, (original_img, original_json, current_img, current_json, status, annotation_count) in enumerate(
self.all_images):
file_item = QTableWidgetItem(os.path.basename(original_img))
status_item = QTableWidgetItem(status)
count_item = QTableWidgetItem(str(annotation_count))
# 设置背景色
if status == "待分类":
color = QColor(255, 200, 200) # 浅红色
elif "err" in status:
color = QColor(255, 102, 102) # 浅红色
else:
color = QColor(200, 255, 200) # 浅绿色
file_item.setBackground(color)
status_item.setBackground(color)
count_item.setBackground(color)
self.table.setItem(i, 0, file_item)
self.table.setItem(i, 1, status_item)
self.table.setItem(i, 2, count_item)
# 显示统计信息
unclassified_count = sum(1 for _, _, _, _, status, _ in self.all_images if status == "待分类")
total_annotations = sum(annotation_count for _, _, _, _, _, annotation_count in self.all_images)
self.statusBar().showMessage(
f"共 {len(self.all_images)} 张图片，{unclassified_count} 张待分类，总标注数: {total_annotations}")
def handle_number_key(self,key_num):
if key_num<len(self.label_map):
self.on_radio_selected(key_num)
def display_image(self, row, column):
"""显示选中的图片及其标注，并更新单选按钮状态"""
if row < 0 or row >= len(self.all_images):
self.image_label.setText("无图片可显示")
self.image_label.setPixmap(QPixmap())
return
original_img, original_json, current_img, current_json, status, annotation_count = self.all_images[row]
# 使用当前路径（如果已移动）或原始路径
img_path = current_img if current_img else original_img
json_path = current_json if current_json else original_json
# 检查文件是否存在
if not os.path.exists(img_path):
QMessageBox.warning(self, "警告", "图片文件不存在！")
self.refresh_image_files()
self.load_image_files()
return
try:
# 加载图片
pixmap = QPixmap(img_path)
if pixmap.isNull():
self.image_label.setText("无法加载图片")
return
# 如果有JSON标注文件，绘制标注框
if json_path and os.path.exists(json_path):
try:
with open(json_path, 'r', encoding='utf-8') as f:
annotation_data = json.load(f)
# 创建带标注的图片
annotated_pixmap = self.draw_annotations(pixmap, annotation_data)
pixmap = annotated_pixmap
except Exception as e:
print(f"加载标注文件失败: {e}")
# 缩放图片以适应显示区域
scaled_pixmap = pixmap.scaled(
self.image_label.width() - 20,
self.image_label.height() - 20,
Qt.KeepAspectRatio,
Qt.SmoothTransformation
)
self.image_label.setPixmap(scaled_pixmap)
# 根据状态设置单选按钮
self.update_radio_buttons(status)
except Exception as e:
self.image_label.setText(f"加载图片出错: {str(e)}")
def update_radio_buttons(self, status_str):
"""根据状态更新单选按钮的选中状态"""
# 清除所有选择
self.radio_group.setExclusive(False)
for btn in self.radio_group.buttons():
btn.setChecked(False)
self.radio_group.setExclusive(True)
self.pic_label.setText(f"图片预览：{status_str}")
# 如果状态是已分类，设置对应的单选按钮
if status_str.startswith("已分类: "):
label_name = status_str.replace("已分类: ", "")
# 找到对应的标签ID
for label_id, name in self.label_map.items():
if name == label_name:
button = self.radio_group.button(label_id)
if button:
button.setChecked(True)
break
def draw_annotations(self, pixmap, annotation_data):
"""在图片上绘制标注框"""
# 创建可绘制的pixmap
result_pixmap = QPixmap(pixmap.size())
result_pixmap.fill(Qt.transparent)
painter = QPainter(result_pixmap)
painter.drawPixmap(0, 0, pixmap)
# 设置画笔
pen = QPen(Qt.red)
pen.setWidth(3)
painter.setPen(pen)
# 绘制每个标注框
for shape in annotation_data.get('shapes', []):
if shape.get('shape_type') == 'rectangle' and len(shape['points']) == 2:
points = shape['points']
x1, y1 = points[0]
x2, y2 = points[1]
# 绘制矩形框
painter.drawRect(int(x1), int(y1), int(x2 - x1), int(y2 - y1))
# 绘制标签
label = shape.get('label', '')
painter.drawText(int(x1), int(y1) - 5, label)
painter.end()
return result_pixmap
if __name__ == "__main__":
base_dir = r"D:\data\course_1027\chan_1028\dan\20251028_1643_part001_seg" # 替换为你的图片目录路径
app = QApplication(sys.argv)
window = ImageAnnotator(base_dir)
window.show()
sys.exit(app.exec_())