375 lines
13 KiB
Python
375 lines
13 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
图像查看器组件模块
|
||
|
||
包含 ImageCategoryTree 和 ImageViewerWidget 类。
|
||
"""
|
||
|
||
import os
|
||
from pathlib import Path
|
||
from typing import Optional, List
|
||
|
||
from PyQt5.QtWidgets import (
|
||
QWidget, QVBoxLayout, QHBoxLayout, QPushButton,
|
||
QFrame, QScrollArea, QLabel, QFileDialog, QMessageBox,
|
||
QTreeWidget, QTreeWidgetItem,
|
||
)
|
||
from PyQt5.QtCore import Qt, QTimer
|
||
from PyQt5.QtGui import QPixmap
|
||
|
||
|
||
class ImageCategoryTree(QTreeWidget):
|
||
"""图像分类目录树 - 按类别组织图像文件"""
|
||
|
||
# 图像类别定义：(类别名称, 关键词列表, 图标)
|
||
CATEGORIES = [
|
||
("模型评估", ["scatter", "regression", "validation", "r2", "rmse"], "📊"),
|
||
("光谱分析", ["spectrum", "spectral", "band", "wavelength"], "📈"),
|
||
("统计图表", ["boxplot", "histogram", "heatmap", "statistics", "stats"], "📉"),
|
||
("处理结果", ["mask", "glint", "deglint", "preview", "overlay", "water_mask"], "🖼️"),
|
||
("含量分布图", [], "📁"),
|
||
]
|
||
|
||
def __init__(self, parent=None):
|
||
super().__init__(parent)
|
||
self.setHeaderLabel("图像目录")
|
||
self.setMaximumWidth(300)
|
||
self.setMinimumWidth(250)
|
||
self.setup_categories()
|
||
self.setStyleSheet("""
|
||
QTreeWidget {
|
||
border: 1px solid #ddd;
|
||
border-radius: 5px;
|
||
background-color: #f8f9fa;
|
||
}
|
||
QTreeWidget::item {
|
||
padding: 5px;
|
||
border-radius: 3px;
|
||
}
|
||
QTreeWidget::item:selected {
|
||
background-color: #0078D4;
|
||
color: white;
|
||
}
|
||
QTreeWidget::item:hover {
|
||
background-color: #e3f2fd;
|
||
}
|
||
""")
|
||
|
||
def setup_categories(self):
|
||
"""初始化类别节点"""
|
||
self.category_items = {}
|
||
for category_name, keywords, icon in self.CATEGORIES:
|
||
item = QTreeWidgetItem(self)
|
||
item.setText(0, f"{icon} {category_name}")
|
||
item.setData(0, Qt.UserRole, {"type": "category", "keywords": keywords, "name": category_name})
|
||
item.setExpanded(True)
|
||
self.category_items[category_name] = item
|
||
|
||
def clear_all_images(self):
|
||
"""清除所有图像项"""
|
||
for category_item in self.category_items.values():
|
||
# 删除所有子项
|
||
while category_item.childCount() > 0:
|
||
category_item.removeChild(category_item.child(0))
|
||
|
||
def add_image(self, file_path: Path, display_name: str = None):
|
||
"""添加图像到对应的类别"""
|
||
if display_name is None:
|
||
display_name = file_path.stem
|
||
|
||
# 根据文件名关键词确定类别
|
||
category = self._determine_category(file_path.name)
|
||
category_item = self.category_items.get(category, self.category_items["含量分布图"])
|
||
|
||
# 创建图像项
|
||
image_item = QTreeWidgetItem(category_item)
|
||
image_item.setText(0, f" └─ {display_name}")
|
||
image_item.setData(0, Qt.UserRole, {"type": "image", "path": str(file_path)})
|
||
image_item.setToolTip(0, str(file_path))
|
||
|
||
return image_item
|
||
|
||
def _determine_category(self, filename: str) -> str:
|
||
"""根据文件名确定类别"""
|
||
filename_lower = filename.lower()
|
||
|
||
for category_name, keywords, _ in self.CATEGORIES:
|
||
if any(keyword in filename_lower for keyword in keywords):
|
||
return category_name
|
||
|
||
return "含量分布图"
|
||
|
||
def scan_directory(self, work_dir: str):
|
||
"""扫描目录中的所有图像文件"""
|
||
self.clear_all_images()
|
||
|
||
work_path = Path(work_dir)
|
||
if not work_path.exists():
|
||
return
|
||
|
||
# 查找所有图像文件：12_visualization 为主，同时扫描步骤产出目录（如 1_water_mask 下的预览/叠置图）
|
||
image_extensions = ['*.png', '*.jpg', '*.jpeg', '*.tif', '*.tiff', '*.bmp']
|
||
scan_roots: List[Path] = []
|
||
_viz = work_path / "12_visualization"
|
||
if _viz.is_dir():
|
||
scan_roots.append(_viz)
|
||
_wm = work_path / "1_water_mask"
|
||
if _wm.is_dir():
|
||
scan_roots.append(_wm)
|
||
if not scan_roots:
|
||
scan_roots.append(work_path)
|
||
|
||
seen_norm: set = set()
|
||
image_files: List[Path] = []
|
||
for root in scan_roots:
|
||
for ext in image_extensions:
|
||
for p in root.glob(f"**/{ext}"):
|
||
key = os.path.normcase(os.path.normpath(str(p.resolve())))
|
||
if key in seen_norm:
|
||
continue
|
||
seen_norm.add(key)
|
||
image_files.append(p)
|
||
|
||
# 添加图像到树
|
||
for img_file in sorted(image_files):
|
||
# 跳过缩略图和临时文件
|
||
if img_file.name.startswith('.') or 'thumb' in img_file.name.lower():
|
||
continue
|
||
self.add_image(img_file)
|
||
|
||
# 更新类别项文本显示数量
|
||
for category_name, item in self.category_items.items():
|
||
count = item.childCount()
|
||
if count > 0:
|
||
for cat_name, _, icon in self.CATEGORIES:
|
||
if cat_name == category_name:
|
||
item.setText(0, f"{icon} {category_name} ({count})")
|
||
break
|
||
|
||
def get_selected_image_path(self) -> Optional[str]:
|
||
"""获取当前选中的图像路径"""
|
||
selected_item = self.currentItem()
|
||
if not selected_item:
|
||
return None
|
||
|
||
data = selected_item.data(0, Qt.UserRole)
|
||
if data and data.get("type") == "image":
|
||
return data.get("path")
|
||
return None
|
||
|
||
|
||
class ImageViewerWidget(QWidget):
|
||
"""图像查看器组件 - 支持缩放、平移"""
|
||
|
||
def __init__(self, parent=None):
|
||
super().__init__(parent)
|
||
self.current_image_path = None
|
||
self.scale_factor = 1.0
|
||
self._update_timer = QTimer() # 防抖定时器
|
||
self._update_timer.setSingleShot(True)
|
||
self._update_timer.timeout.connect(self._do_update_display)
|
||
self._pending_scale = None # 待更新的缩放比例
|
||
self.setup_ui()
|
||
|
||
def setup_ui(self):
|
||
layout = QVBoxLayout()
|
||
layout.setContentsMargins(0, 0, 0, 0)
|
||
|
||
# 工具栏
|
||
toolbar = QHBoxLayout()
|
||
|
||
self.refresh_btn = QPushButton("🔄 刷新目录")
|
||
self.refresh_btn.setToolTip("重新扫描工作目录中的图像文件")
|
||
toolbar.addWidget(self.refresh_btn)
|
||
|
||
# 添加分隔线
|
||
separator = QFrame()
|
||
separator.setFrameShape(QFrame.VLine)
|
||
separator.setFrameShadow(QFrame.Sunken)
|
||
toolbar.addWidget(separator)
|
||
|
||
self.zoom_in_btn = QPushButton("🔍+")
|
||
self.zoom_in_btn.setToolTip("放大")
|
||
self.zoom_in_btn.setMaximumWidth(50)
|
||
toolbar.addWidget(self.zoom_in_btn)
|
||
|
||
self.zoom_out_btn = QPushButton("🔍-")
|
||
self.zoom_out_btn.setToolTip("缩小")
|
||
self.zoom_out_btn.setMaximumWidth(50)
|
||
toolbar.addWidget(self.zoom_out_btn)
|
||
|
||
self.fit_btn = QPushButton("⬜ 适应窗口")
|
||
self.fit_btn.setToolTip("适应窗口大小")
|
||
toolbar.addWidget(self.fit_btn)
|
||
|
||
self.original_btn = QPushButton("1:1 原始大小")
|
||
self.original_btn.setToolTip("原始大小")
|
||
toolbar.addWidget(self.original_btn)
|
||
|
||
toolbar.addStretch()
|
||
|
||
self.save_btn = QPushButton("💾 保存")
|
||
self.save_btn.setToolTip("保存当前图像")
|
||
toolbar.addWidget(self.save_btn)
|
||
|
||
layout.addLayout(toolbar)
|
||
|
||
# 图像显示区域 - 使用 QLabel + QScrollArea
|
||
self.scroll_area = QScrollArea()
|
||
self.scroll_area.setWidgetResizable(True)
|
||
self.scroll_area.setStyleSheet("background-color: white;")
|
||
|
||
self.image_label = QLabel()
|
||
self.image_label.setAlignment(Qt.AlignCenter)
|
||
self.image_label.setStyleSheet("background-color: white;")
|
||
|
||
self.scroll_area.setWidget(self.image_label)
|
||
layout.addWidget(self.scroll_area, 1)
|
||
|
||
# 状态栏
|
||
status_layout = QHBoxLayout()
|
||
self.status_label = QLabel("就绪")
|
||
self.status_label.setStyleSheet("color: #666; font-size: 11px;")
|
||
status_layout.addWidget(self.status_label)
|
||
status_layout.addStretch()
|
||
layout.addLayout(status_layout)
|
||
|
||
self.setLayout(layout)
|
||
|
||
# 连接信号
|
||
self.zoom_in_btn.clicked.connect(self.zoom_in)
|
||
self.zoom_out_btn.clicked.connect(self.zoom_out)
|
||
self.fit_btn.clicked.connect(self.fit_to_window)
|
||
self.original_btn.clicked.connect(self.original_size)
|
||
self.save_btn.clicked.connect(self.save_image)
|
||
|
||
def load_image(self, image_path: str):
|
||
"""加载并显示图像"""
|
||
if not image_path or not Path(image_path).exists():
|
||
self.image_label.setText("图像不存在")
|
||
self.status_label.setText("图像加载失败")
|
||
return
|
||
|
||
self.current_image_path = image_path
|
||
self.scale_factor = 1.0
|
||
|
||
# 加载图像
|
||
pixmap = QPixmap(image_path)
|
||
if pixmap.isNull():
|
||
self.image_label.setText("无法加载图像")
|
||
self.status_label.setText("图像格式不支持")
|
||
return
|
||
|
||
self.original_pixmap = pixmap
|
||
|
||
# 默认适应窗口显示
|
||
self.fit_to_window()
|
||
|
||
# 更新状态
|
||
file_info = Path(image_path).stat()
|
||
size_mb = file_info.st_size / (1024 * 1024)
|
||
self.status_label.setText(f"{pixmap.width()}x{pixmap.height()} | {size_mb:.2f} MB | {Path(image_path).name} | 适应窗口")
|
||
|
||
def update_image_display(self):
|
||
"""更新图像显示 - 使用防抖避免频繁重绘卡顿"""
|
||
# 取消之前的待执行更新，重新计时
|
||
self._update_timer.stop()
|
||
self._pending_scale = self.scale_factor
|
||
self._update_timer.start(50) # 50ms后执行实际更新
|
||
|
||
def _do_update_display(self):
|
||
"""实际执行图像更新"""
|
||
if not hasattr(self, 'original_pixmap') or self.original_pixmap.isNull():
|
||
return
|
||
|
||
if self._pending_scale is None:
|
||
return
|
||
|
||
# 根据缩放比例选择变换模式：大幅度缩放用Fast模式提升性能
|
||
if self._pending_scale > 2.0 or self._pending_scale < 0.5:
|
||
transform = Qt.FastTransformation
|
||
else:
|
||
transform = Qt.SmoothTransformation
|
||
|
||
scaled_pixmap = self.original_pixmap.scaled(
|
||
int(self.original_pixmap.width() * self._pending_scale),
|
||
int(self.original_pixmap.height() * self._pending_scale),
|
||
Qt.KeepAspectRatio,
|
||
transform
|
||
)
|
||
self.image_label.setPixmap(scaled_pixmap)
|
||
self._pending_scale = None
|
||
|
||
def wheelEvent(self, event):
|
||
"""鼠标滚轮缩放 - 实时响应"""
|
||
delta = event.angleDelta().y()
|
||
|
||
if delta > 0:
|
||
# 向上滚动 - 放大
|
||
if self.scale_factor < 5.0:
|
||
self.scale_factor = min(self.scale_factor * 1.1, 5.0)
|
||
self.update_image_display()
|
||
else:
|
||
# 向下滚动 - 缩小
|
||
if self.scale_factor > 0.1:
|
||
self.scale_factor = max(self.scale_factor / 1.1, 0.1)
|
||
self.update_image_display()
|
||
|
||
event.accept()
|
||
|
||
def zoom_in(self):
|
||
"""放大"""
|
||
if self.scale_factor < 5.0:
|
||
self.scale_factor = min(self.scale_factor * 1.25, 5.0)
|
||
self.update_image_display()
|
||
|
||
def zoom_out(self):
|
||
"""缩小"""
|
||
if self.scale_factor > 0.1:
|
||
self.scale_factor = max(self.scale_factor / 1.25, 0.1)
|
||
self.update_image_display()
|
||
|
||
def fit_to_window(self):
|
||
"""适应窗口"""
|
||
if not hasattr(self, 'original_pixmap') or self.original_pixmap.isNull():
|
||
return
|
||
|
||
# 计算适应窗口的缩放比例
|
||
view_size = self.scroll_area.viewport().size()
|
||
img_size = self.original_pixmap.size()
|
||
|
||
scale_w = view_size.width() / img_size.width()
|
||
scale_h = view_size.height() / img_size.height()
|
||
|
||
# 记录适应前的比例（用于后续恢复参考）
|
||
self._fit_scale = min(scale_w, scale_h)
|
||
self.scale_factor = self._fit_scale
|
||
|
||
self.update_image_display()
|
||
self.status_label.setText(f"适应窗口 | 缩放: {self.scale_factor:.1%}")
|
||
|
||
def original_size(self):
|
||
"""原始大小"""
|
||
self.scale_factor = 1.0
|
||
self._fit_scale = None # 清除适应记录
|
||
self.update_image_display()
|
||
self.status_label.setText("原始大小 | 缩放: 100%")
|
||
|
||
def save_image(self):
|
||
"""保存图像"""
|
||
if not self.current_image_path:
|
||
return
|
||
|
||
file_path, _ = QFileDialog.getSaveFileName(
|
||
self, "保存图像", Path(self.current_image_path).name,
|
||
"PNG图片 (*.png);;JPG图片 (*.jpg);;所有文件 (*.*)"
|
||
)
|
||
|
||
if file_path:
|
||
try:
|
||
import shutil
|
||
shutil.copy(self.current_image_path, file_path)
|
||
except Exception as e:
|
||
QMessageBox.critical(self, "错误", f"保存失败: {e}")
|