全程键盘标注，方便快捷
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 natsortedclass CustomTableWidget(QTableWidget):def __init__(self, parent=None):super().__init__(parent)self.parent_window = parentdef 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 - 1self.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 + 1self.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_0if 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_0print(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_dirself.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_pathdef 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 = pathelse: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] ifj.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_pathreturn Nonedef get_annotation_count(self, json_path):"""从JSON文件中获取标注数量"""if not json_path or not os.path.exists(json_path):return 0try:with open(json_path, 'r', encoding='utf-8') as f:data = json.load(f)return len(data.get('shapes', []))except:return 0def move_image_and_json(self, img_path, json_path, label_id):"""将图片和JSON文件移动到对应的分类目录"""try:if label_id not in self.output_dirs:return None, Nonetarget_dir = self.output_dirs[label_id]img_filename = os.path.basename(img_path)# 移动图片文件img_target_path = os.path.join(target_dir, img_filename)counter = 1base_name, ext = os.path.splitext(img_filename)if os.path.exists(img_target_path):json_target_path = Noneif 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_pathshutil.move(img_path, img_target_path)# 移动JSON文件（如果存在）json_target_path = Noneif 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_pathexcept Exception as e:print(f"移动文件失败: {e}")return None, Nonedef 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_imgjson_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+1next_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 ireturn -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())returnoriginal_img, original_json, current_img, current_json, status, annotation_count = self.all_images[row]# 使用当前路径（如果已移动）或原始路径img_path = current_img if current_img else original_imgjson_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()returntry:# 加载图片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_pixmapexcept 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("已分类: ", "")# 找到对应的标签IDfor label_id, name in self.label_map.items():if name == label_name:button = self.radio_group.button(label_id)if button:button.setChecked(True)breakdef draw_annotations(self, pixmap, annotation_data):"""在图片上绘制标注框"""# 创建可绘制的pixmapresult_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_pixmapif __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_())
