Commit 08cb52c5 by Yuhaibo

1

parent 0e36f2c7
...@@ -54,7 +54,7 @@ class ModelSetHandler: ...@@ -54,7 +54,7 @@ class ModelSetHandler:
type_layout = QtWidgets.QHBoxLayout() type_layout = QtWidgets.QHBoxLayout()
type_layout.addWidget(QtWidgets.QLabel("模型类型:")) type_layout.addWidget(QtWidgets.QLabel("模型类型:"))
type_combo = QtWidgets.QComboBox() type_combo = QtWidgets.QComboBox()
type_combo.addItems(["YOLOv8", "YOLOv11", "Faster R-CNN", "SSD", "RetinaNet", "自定义"])
type_layout.addWidget(type_combo) type_layout.addWidget(type_combo)
layout.addLayout(type_layout) layout.addLayout(type_layout)
......
# AutoAnnotationDetector 使用说明
## 功能概述
`AutoAnnotationDetector` 整合了标框和标点两个功能,共享相同的检测逻辑,但提供不同的输出方式。
## 核心设计
```
输入图像 → detect() → 检测结果
┌─────────┴─────────┐
↓ ↓
get_boxes() get_points()
↓ ↓
box位置数据 点位置数据
↓ ↓
draw_boxes() draw_points()
↓ ↓
绘制框的图像 绘制点的图像
```
## 快速开始
```python
from handlers.videopage.auto_dot import AutoAnnotationDetector
import cv2
# 1. 创建检测器
detector = AutoAnnotationDetector(
model_path="path/to/model.dat",
device='cuda'
)
# 2. 加载图像
image = cv2.imread("test.jpg")
# 3. 执行核心检测(只需执行一次)
detection_result = detector.detect(
image=image,
conf_threshold=0.5,
min_area=100
)
# 4a. 获取box位置数据
boxes = detector.get_boxes(
detection_result,
padding=10,
merge_all=True # True=合并所有类别, False=按类别分别生成
)
# 4b. 获取点位置数据
points = detector.get_points(detection_result)
# 5a. 绘制box
box_image = detector.draw_boxes(image, boxes)
cv2.imwrite("box_result.jpg", box_image)
# 5b. 绘制点
point_image = detector.draw_points(image, points)
cv2.imwrite("point_result.jpg", point_image)
```
## API 详解
### 1. detect() - 核心检测方法
执行YOLO推理,返回分割结果。
**参数:**
- `image`: 输入图像 (numpy.ndarray)
- `conf_threshold`: 置信度阈值 (默认0.5)
- `min_area`: 最小面积阈值,过滤小区域 (默认100像素)
**返回:**
```python
{
'success': bool,
'masks': List[np.ndarray], # 有效的mask列表
'class_names': List[str], # 对应的类别名称
'confidences': List[float], # 对应的置信度
'image_shape': tuple # 图像尺寸 (height, width)
}
```
### 2. get_boxes() - 生成box位置数据
从检测结果生成box坐标。
**参数:**
- `detection_result`: detect()的返回结果
- `padding`: box边界扩展像素数 (默认10)
- `merge_all`: 是否合并所有类别 (默认True)
**返回:**
```python
[
{
'x1': int, # 左上角x
'y1': int, # 左上角y
'x2': int, # 右下角x
'y2': int, # 右下角y
'classes': List[str], # 包含的类别
'area': int # box面积
},
...
]
```
### 3. get_points() - 生成点位置数据
从检测结果生成容器顶部和底部点坐标。
**参数:**
- `detection_result`: detect()的返回结果
**返回:**
```python
[
{
'top': int, # 顶部y坐标
'bottom': int, # 底部y坐标
'top_x': int, # 顶部x坐标
'bottom_x': int, # 底部x坐标
'height': int, # 容器高度
'method': str, # 检测方法
'confidence': float # 置信度
},
...
]
```
### 4. draw_boxes() - 绘制box
在图像上绘制box框。
**参数:**
- `image`: 输入图像
- `boxes`: get_boxes()返回的box列表
**返回:** 绘制完成的图像
### 5. draw_points() - 绘制点
在图像上绘制顶部和底部点。
**参数:**
- `image`: 输入图像
- `points`: get_points()返回的点列表
**返回:** 绘制完成的图像
## 使用场景
### 场景1: 只需要box标注
```python
detection_result = detector.detect(image)
boxes = detector.get_boxes(detection_result, merge_all=True)
annotated_image = detector.draw_boxes(image, boxes)
```
### 场景2: 只需要点标注
```python
detection_result = detector.detect(image)
points = detector.get_points(detection_result)
annotated_image = detector.draw_points(image, points)
```
### 场景3: 同时需要box和点
```python
# 只需执行一次检测
detection_result = detector.detect(image)
# 分别获取box和点
boxes = detector.get_boxes(detection_result)
points = detector.get_points(detection_result)
# 分别绘制
box_image = detector.draw_boxes(image, boxes)
point_image = detector.draw_points(image, points)
```
### 场景4: 获取位置数据但不绘制
```python
detection_result = detector.detect(image)
boxes = detector.get_boxes(detection_result)
points = detector.get_points(detection_result)
# 直接使用位置数据,不绘制图像
for box in boxes:
print(f"Box: ({box['x1']}, {box['y1']}) -> ({box['x2']}, {box['y2']})")
for point in points:
print(f"Top: ({point['top_x']}, {point['top']})")
print(f"Bottom: ({point['bottom_x']}, {point['bottom']})")
```
## 点检测逻辑
根据检测到的类别组合,自动选择最佳检测方法:
1. **liquid_only**: 仅liquid → liquid最低点 + liquid最高点
2. **air_only**: 仅air → air最低点 + air最高点
3. **foam_only**: 仅foam → foam最低点 + foam最高点
4. **liquid_air**: liquid + air → liquid最低点 + air最高点
5. **liquid_foam**: liquid + foam → liquid最低点 + foam最高点
6. **liquid_foam_air**: liquid + foam + air → liquid最低点 + air最高点
7. **foam_air**: foam + air → foam最低点 + air最高点
## 兼容性
为保持向后兼容,提供了别名:
```python
AutoboxDetector = AutoAnnotationDetector
```
旧代码可以继续使用 `AutoboxDetector`
## 测试
运行测试函数:
```bash
python auto_dot.py
```
测试会自动:
1. 执行一次检测
2. 生成box标注并保存
3. 生成点标注并保存
4. 显示结果图像
# 自动标点功能模块使用说明 # 自动标点与自动选框功能模块使用说明
## 功能概述 ## 功能概述
`auto_dot.py` 模块实现了基于YOLO分割掩码的自动标点功能,可以自动检测容器的顶部和底部位置,替代人工手动标点。 `auto_dot.py` 模块实现了两个核心功能:
1. **AutoDotDetector**: 基于YOLO分割掩码的自动标点功能,自动检测容器的顶部和底部位置
2. **AutoboxDetector**: 自动选择标注区域功能,自动生成包含分割结果的最小面积box
## 核心特性 ## 核心特性
...@@ -153,14 +155,109 @@ if result['success']: ...@@ -153,14 +155,109 @@ if result['success']:
3. **调整检测框**: 如果检测失败,尝试调整检测框的位置和大小 3. **调整检测框**: 如果检测失败,尝试调整检测框的位置和大小
4. **降低置信度**: 如果检测不到掩码,尝试降低 `conf_threshold` 4. **降低置信度**: 如果检测不到掩码,尝试降低 `conf_threshold`
## AutoboxDetector - 自动选框功能
### 功能说明
自动选择标注区域类,输入图像后自动生成包含分割结果的最小面积box。
### 核心特性
- **输入**: 图像
- **输出**: box位置数据 + 绘制完框的图像
- **box设置逻辑**:
1. 对每个类别(liquid、air、foam)的分割结果,合并同类别的所有mask
2. 计算每个类别mask的最小包围框
3. 支持设置最小面积阈值过滤小区域
4. 支持设置padding扩展box边界
### API 使用示例
```python
from handlers.videopage.auto_dot import AutoboxDetector
import cv2
# 1. 创建检测器
detector = AutoboxDetector(
model_path="path/to/model.dat",
device='cuda' # 或 'cpu'
)
# 2. 加载图片
image = cv2.imread("test_image.jpg")
# 3. 执行检测
result = detector.detect_boxes(
image=image,
conf_threshold=0.5, # 置信度阈值
min_area=100, # 最小面积阈值(像素)
padding=10 # box边界扩展像素数
)
# 4. 获取结果
if result['success']:
for box in result['boxes']:
print(f"类别: {box['class']}")
print(f"坐标: ({box['x1']}, {box['y1']}) -> ({box['x2']}, {box['y2']})")
print(f"面积: {box['area']}px²")
# 保存标注图片
cv2.imwrite("box_result.jpg", result['annotated_image'])
```
### 输出数据结构
```python
{
'success': bool, # 检测是否成功
'boxes': [
{
'x1': int, # 左上角x坐标
'y1': int, # 左上角y坐标
'x2': int, # 右下角x坐标
'y2': int, # 右下角y坐标
'class': str, # 类别名称 (liquid/air/foam)
'area': int # box面积(像素²)
},
...
],
'annotated_image': np.ndarray # 绘制完框的图像
}
```
### 可视化标注
标注图片包含:
- **绿色框**: liquid区域
- **蓝色框**: air区域
- **橙色框**: foam区域
- **类别标签**: 显示在框的左上角
### 独立测试
运行测试函数:
```bash
cd D:\restructure\liquid_level_line_detection_system\handlers\videopage
python auto_dot.py
```
测试函数会自动调用 `test_auto_box()`,输出结果保存在:
```
D:\restructure\liquid_level_line_detection_system\test_output\auto_box_result.jpg
```
## 接入系统 ## 接入系统
调试成功后,可以在主系统中调用: 调试成功后,可以在主系统中调用:
```python ```python
from handlers.videopage.auto_dot import AutoDotDetector from handlers.videopage.auto_dot import AutoDotDetector, AutoboxDetector
# 方式1: 自动标点 - 在标注页面添加"自动标点"按钮
detector = AutoDotDetector(model_path="path/to/model.dat")
result = detector.detect_container_boundaries(image, conf_threshold=0.5)
# 在标注页面添加"自动标点"按钮 # 方式2: 自动选框 - 在标注页面添加"自动选框"按钮
# 点击后调用 detector.detect_container_boundaries() detector = AutoboxDetector(model_path="path/to/model.dat")
# 将返回的 top/bottom 坐标填充到标注点位置 result = detector.detect_boxes(image, conf_threshold=0.5, min_area=100, padding=10)
``` ```
...@@ -751,5 +751,953 @@ def test_auto_dot(): ...@@ -751,5 +751,953 @@ def test_auto_dot():
print(f"{'='*80}") print(f"{'='*80}")
class AutoAnnotationDetector:
"""
自动标注检测器(整合标框和标点功能)
功能:
1. 输入图像
2. 使用YOLO分割模型检测液体、空气、泡沫区域
3. 可输出box位置数据(标框)或点位置数据(标点)
4. 根据位置信息绘制框或点
"""
def __init__(self, model_path: str = None, device: str = 'cuda'):
"""
初始化自动标注检测器
Args:
model_path: YOLO模型路径(.pt 或 .dat 文件)
device: 计算设备 ('cuda' 或 'cpu')
"""
self.model = None
self.model_path = model_path
self.device = self._validate_device(device)
if model_path:
self.load_model(model_path)
def _validate_device(self, device: str) -> str:
"""验证并选择可用的设备"""
try:
import torch
if device in ['cuda', '0'] or device.startswith('cuda:'):
if torch.cuda.is_available():
return 'cuda' if device in ['cuda', '0'] else device
else:
print("⚠️ CUDA不可用,切换到CPU")
return 'cpu'
return device
except Exception:
return 'cpu'
def load_model(self, model_path: str) -> bool:
"""
加载YOLO模型
Args:
model_path: 模型文件路径
Returns:
bool: 加载是否成功
"""
try:
import os
if not os.path.exists(model_path):
print(f"❌ 模型文件不存在: {model_path}")
return False
# 如果是 .dat 文件,先解码
if model_path.endswith('.dat'):
decoded_path = self._decode_dat_model(model_path)
if decoded_path is None:
print(f"❌ .dat 文件解码失败: {model_path}")
return False
model_path = decoded_path
# 设置离线模式
os.environ['YOLO_VERBOSE'] = 'False'
os.environ['YOLO_OFFLINE'] = '1'
os.environ['ULTRALYTICS_OFFLINE'] = 'True'
from ultralytics import YOLO
print(f"🔄 正在加载模型: {model_path}")
self.model = YOLO(model_path)
self.model.to(self.device)
self.model_path = model_path
print(f"✅ 模型加载成功: {os.path.basename(model_path)}")
return True
except Exception as e:
print(f"❌ 模型加载失败: {e}")
return False
def _decode_dat_model(self, dat_path: str) -> Optional[str]:
"""解码 .dat 格式的模型文件"""
try:
import struct
import hashlib
SIGNATURE = b'LDS_MODEL_FILE'
VERSION = 1
ENCRYPTION_KEY = "liquid_detection_system_2024"
key_hash = hashlib.sha256(ENCRYPTION_KEY.encode('utf-8')).digest()
with open(dat_path, 'rb') as f:
signature = f.read(len(SIGNATURE))
if signature != SIGNATURE:
return None
version = struct.unpack('<I', f.read(4))[0]
if version != VERSION:
return None
filename_len = struct.unpack('<I', f.read(4))[0]
original_filename = f.read(filename_len).decode('utf-8')
data_len = struct.unpack('<Q', f.read(8))[0]
encrypted_data = f.read(data_len)
# XOR 解密
decrypted_data = bytearray()
key_len = len(key_hash)
for i, byte in enumerate(encrypted_data):
decrypted_data.append(byte ^ key_hash[i % key_len])
decrypted_data = bytes(decrypted_data)
# 保存到临时目录
temp_dir = Path(get_temp_models_dir())
temp_dir.mkdir(exist_ok=True)
path_hash = hashlib.md5(str(dat_path).encode()).hexdigest()[:8]
temp_model_path = temp_dir / f"temp_{Path(dat_path).stem}_{path_hash}.pt"
with open(temp_model_path, 'wb') as f:
f.write(decrypted_data)
return str(temp_model_path)
except Exception as e:
print(f"❌ 解码.dat文件失败: {e}")
return None
def detect(
self,
image: np.ndarray,
conf_threshold: float = 0.5,
min_area: int = 100
) -> Dict:
"""
核心检测方法 - 执行YOLO推理并返回分割结果
Args:
image: 输入图片 (numpy.ndarray, BGR格式)
conf_threshold: 置信度阈值
min_area: 最小面积阈值(像素),小于此面积的mask会被过滤
Returns:
dict: {
'success': bool,
'masks': List[np.ndarray], # 有效的mask列表
'class_names': List[str], # 对应的类别名称
'confidences': List[float], # 对应的置信度
'image_shape': tuple # 图像尺寸 (height, width)
}
"""
if self.model is None:
return {
'success': False,
'boxes': [],
'annotated_image': image.copy(),
'error': '模型未加载'
}
if image.size == 0:
return {
'success': False,
'boxes': [],
'annotated_image': image.copy(),
'error': '图像为空'
}
print(f"\n{'='*60}")
print(f"🔍 开始自动检测")
print(f" 图像尺寸: {image.shape[1]}x{image.shape[0]}")
# 执行YOLO推理
print(f" 🔄 执行YOLO推理...")
try:
mission_results = self.model.predict(
source=image,
imgsz=640,
conf=conf_threshold,
iou=0.5,
device=self.device,
save=False,
verbose=False,
half=True if self.device != 'cpu' else False,
stream=False
)
mission_result = mission_results[0]
except Exception as e:
print(f" ❌ YOLO推理失败: {e}")
return {
'success': False,
'boxes': [],
'annotated_image': image.copy(),
'error': f'YOLO推理失败: {e}'
}
# 处理检测结果
if mission_result.masks is None:
print(f" ⚠️ 未检测到任何掩码")
return {
'success': False,
'boxes': [],
'annotated_image': image.copy(),
'error': '未检测到任何掩码'
}
masks = mission_result.masks.data.cpu().numpy() > 0.5
classes = mission_result.boxes.cls.cpu().numpy().astype(int)
confidences = mission_result.boxes.conf.cpu().numpy()
print(f" ✅ 检测到 {len(masks)} 个对象")
# 收集所有有效的mask
valid_masks = []
class_names = []
for i in range(len(masks)):
if confidences[i] < conf_threshold:
continue
class_name = self.model.names[classes[i]]
conf = confidences[i]
# 调整mask尺寸到输入图像大小
resized_mask = cv2.resize(
masks[i].astype(np.uint8),
(image.shape[1], image.shape[0])
) > 0.5
# 计算mask面积
mask_area = np.sum(resized_mask)
print(f" - {class_name}: {mask_area} 像素, 置信度: {conf:.3f}")
# 过滤小面积mask
if mask_area < min_area:
print(f" ⚠️ 面积过小,已过滤")
continue
valid_masks.append(resized_mask)
class_names.append(class_name)
if len(valid_masks) == 0:
print(f" ⚠️ 没有有效的mask")
return {
'success': False,
'error': '没有有效的mask'
}
print(f"\n{'='*60}")
print(f"✅ 检测完成,共 {len(valid_masks)} 个有效mask")
return {
'success': True,
'masks': valid_masks,
'class_names': class_names,
'confidences': [confidences[i] for i in range(len(masks)) if i < len(valid_masks)],
'image_shape': (image.shape[0], image.shape[1])
}
def get_boxes(
self,
detection_result: Dict,
padding: int = 10,
merge_all: bool = True
) -> List[Dict]:
"""
从检测结果生成box位置数据
Args:
detection_result: detect()方法的返回结果
padding: box边界扩展像素数
merge_all: 是否合并所有类别生成一个大框
Returns:
List[Dict]: box列表,每个box包含 {'x1', 'y1', 'x2', 'y2', 'classes', 'area'}
"""
if not detection_result.get('success'):
return []
valid_masks = detection_result['masks']
class_names = detection_result['class_names']
height, width = detection_result['image_shape']
boxes = []
if merge_all:
# 合并所有mask生成一个大框
combined_mask = np.zeros_like(valid_masks[0], dtype=bool)
for mask in valid_masks:
combined_mask = combined_mask | mask
# 找到mask的边界
y_coords, x_coords = np.where(combined_mask)
if len(y_coords) > 0 and len(x_coords) > 0:
# 计算最小包围框
x1 = max(0, int(np.min(x_coords)) - padding)
y1 = max(0, int(np.min(y_coords)) - padding)
x2 = min(width, int(np.max(x_coords)) + padding)
y2 = min(height, int(np.max(y_coords)) + padding)
box_area = (x2 - x1) * (y2 - y1)
# 统计包含的类别
unique_classes = list(set(class_names))
boxes.append({
'x1': x1,
'y1': y1,
'x2': x2,
'y2': y2,
'classes': unique_classes, # 包含的所有类别
'area': box_area
})
print(f" 📦 生成合并box: ({x1}, {y1}) -> ({x2}, {y2})")
print(f" 包含类别: {', '.join(unique_classes)}")
print(f" 面积: {box_area}px²")
else:
# 按类别分别生成框
class_masks = {}
for mask, class_name in zip(valid_masks, class_names):
if class_name not in class_masks:
class_masks[class_name] = []
class_masks[class_name].append(mask)
for class_name, masks_list in class_masks.items():
# 合并同类别的所有mask
combined_mask = np.zeros_like(masks_list[0], dtype=bool)
for mask in masks_list:
combined_mask = combined_mask | mask
# 找到mask的边界
y_coords, x_coords = np.where(combined_mask)
if len(y_coords) == 0 or len(x_coords) == 0:
continue
# 计算最小包围框
x1 = max(0, int(np.min(x_coords)) - padding)
y1 = max(0, int(np.min(y_coords)) - padding)
x2 = min(width, int(np.max(x_coords)) + padding)
y2 = min(height, int(np.max(y_coords)) + padding)
box_area = (x2 - x1) * (y2 - y1)
boxes.append({
'x1': x1,
'y1': y1,
'x2': x2,
'y2': y2,
'classes': [class_name],
'area': box_area
})
print(f" 📦 生成box [{class_name}]: ({x1}, {y1}) -> ({x2}, {y2}), 面积={box_area}px²")
return boxes
def get_points(
self,
detection_result: Dict
) -> List[Dict]:
"""
从检测结果生成点位置数据(容器顶部和底部)
Args:
detection_result: detect()方法的返回结果
Returns:
List[Dict]: 点列表,每个点包含 {'top', 'bottom', 'top_x', 'bottom_x', 'height', 'method', 'classes'}
"""
if not detection_result.get('success'):
return []
valid_masks = detection_result['masks']
class_names = detection_result['class_names']
# 按类别收集坐标
liquid_y_coords = []
liquid_x_coords = []
air_y_coords = []
air_x_coords = []
foam_y_coords = []
foam_x_coords = []
for mask, class_name in zip(valid_masks, class_names):
y_coords, x_coords = np.where(mask)
if class_name == 'liquid':
liquid_y_coords.extend(y_coords)
liquid_x_coords.extend(x_coords)
elif class_name == 'air':
air_y_coords.extend(y_coords)
air_x_coords.extend(x_coords)
elif class_name == 'foam':
foam_y_coords.extend(y_coords)
foam_x_coords.extend(x_coords)
# 使用AutoDotDetector的逻辑计算边界
container_info = self._calculate_boundaries(
liquid_y_coords, liquid_x_coords,
air_y_coords, air_x_coords,
[(None, foam_y_coords, foam_x_coords)] if len(foam_y_coords) > 0 else [],
1.0, 1.0, 1.0,
0, 0, 0
)
if container_info:
return [container_info]
else:
return []
def draw_boxes(
self,
image: np.ndarray,
boxes: List[Dict]
) -> np.ndarray:
"""
在图像上绘制box
Args:
image: 输入图像
boxes: box列表
Returns:
绘制完成的图像
"""
annotated_image = image.copy()
self._draw_boxes(annotated_image, boxes)
return annotated_image
def draw_points(
self,
image: np.ndarray,
points: List[Dict]
) -> np.ndarray:
"""
在图像上绘制点
Args:
image: 输入图像
points: 点列表
Returns:
绘制完成的图像
"""
annotated_image = image.copy()
for point_info in points:
self._draw_annotations(
annotated_image,
point_info,
0, image.shape[1]
)
return annotated_image
def draw_all(
self,
image: np.ndarray,
boxes: List[Dict],
points: List[Dict]
) -> np.ndarray:
"""
在同一张图像上同时绘制框和点
Args:
image: 输入图像
boxes: box列表
points: 点列表
Returns:
绘制完成的图像
"""
annotated_image = image.copy()
# 先绘制框
self._draw_boxes(annotated_image, boxes)
# 再绘制点
for point_info in points:
self._draw_annotations(
annotated_image,
point_info,
0, image.shape[1]
)
return annotated_image
def _calculate_boundaries(
self,
liquid_y_coords: List[int],
liquid_x_coords: List[int],
air_y_coords: List[int],
air_x_coords: List[int],
foam_masks: List,
liquid_conf: float,
air_conf: float,
foam_conf: float,
crop_top: int,
crop_left: int,
idx: int
) -> Optional[Dict]:
"""计算容器的顶部和底部边界(复用AutoDotDetector的逻辑)"""
container_bottom = None
container_top = None
bottom_x = None
top_x = None
method = None
confidence = 0.0
has_liquid = len(liquid_y_coords) > 0
has_air = len(air_y_coords) > 0
has_foam = len(foam_masks) > 0
# 情况5: liquid + foam + air
if has_liquid and has_foam and has_air:
liquid_y_array = np.array(liquid_y_coords)
liquid_x_array = np.array(liquid_x_coords)
max_y = np.max(liquid_y_array)
bottom_region_mask = liquid_y_array >= (max_y - 5)
container_bottom_in_crop = int(np.median(liquid_y_array[bottom_region_mask]))
on_bottom_line = liquid_y_array == container_bottom_in_crop
if np.sum(on_bottom_line) > 0:
bottom_x_in_crop = int(np.median(liquid_x_array[on_bottom_line]))
else:
bottom_x_in_crop = int(np.median(liquid_x_array[bottom_region_mask]))
air_y_array = np.array(air_y_coords)
air_x_array = np.array(air_x_coords)
min_y = np.min(air_y_array)
top_region_mask = air_y_array <= (min_y + 5)
container_top_in_crop = int(np.median(air_y_array[top_region_mask]))
on_top_line = air_y_array == container_top_in_crop
if np.sum(on_top_line) > 0:
top_x_in_crop = int(np.median(air_x_array[on_top_line]))
else:
top_x_in_crop = int(np.median(air_x_array[top_region_mask]))
container_bottom = container_bottom_in_crop + crop_top
container_top = container_top_in_crop + crop_top
bottom_x = bottom_x_in_crop + crop_left
top_x = top_x_in_crop + crop_left
method = 'liquid_foam_air'
confidence = (liquid_conf + air_conf + foam_conf) / 3
# 情况4: liquid + air
elif has_liquid and has_air:
liquid_y_array = np.array(liquid_y_coords)
liquid_x_array = np.array(liquid_x_coords)
max_y = np.max(liquid_y_array)
bottom_region_mask = liquid_y_array >= (max_y - 5)
container_bottom_in_crop = int(np.median(liquid_y_array[bottom_region_mask]))
on_bottom_line = liquid_y_array == container_bottom_in_crop
if np.sum(on_bottom_line) > 0:
bottom_x_in_crop = int(np.median(liquid_x_array[on_bottom_line]))
else:
bottom_x_in_crop = int(np.median(liquid_x_array[bottom_region_mask]))
air_y_array = np.array(air_y_coords)
air_x_array = np.array(air_x_coords)
min_y = np.min(air_y_array)
top_region_mask = air_y_array <= (min_y + 5)
container_top_in_crop = int(np.median(air_y_array[top_region_mask]))
on_top_line = air_y_array == container_top_in_crop
if np.sum(on_top_line) > 0:
top_x_in_crop = int(np.median(air_x_array[on_top_line]))
else:
top_x_in_crop = int(np.median(air_x_array[top_region_mask]))
container_bottom = container_bottom_in_crop + crop_top
container_top = container_top_in_crop + crop_top
bottom_x = bottom_x_in_crop + crop_left
top_x = top_x_in_crop + crop_left
method = 'liquid_air'
confidence = (liquid_conf + air_conf) / 2
# 情况6: liquid + foam
elif has_liquid and has_foam:
liquid_y_array = np.array(liquid_y_coords)
liquid_x_array = np.array(liquid_x_coords)
max_y = np.max(liquid_y_array)
bottom_region_mask = liquid_y_array >= (max_y - 5)
container_bottom_in_crop = int(np.median(liquid_y_array[bottom_region_mask]))
on_bottom_line = liquid_y_array == container_bottom_in_crop
if np.sum(on_bottom_line) > 0:
bottom_x_in_crop = int(np.median(liquid_x_array[on_bottom_line]))
else:
bottom_x_in_crop = int(np.median(liquid_x_array[bottom_region_mask]))
all_foam_y = []
all_foam_x = []
for mask, y_coords, x_coords in foam_masks:
all_foam_y.extend(y_coords)
all_foam_x.extend(x_coords)
foam_y_array = np.array(all_foam_y)
foam_x_array = np.array(all_foam_x)
min_y = np.min(foam_y_array)
top_region_mask = foam_y_array <= (min_y + 5)
container_top_in_crop = int(np.median(foam_y_array[top_region_mask]))
on_top_line = foam_y_array == container_top_in_crop
if np.sum(on_top_line) > 0:
top_x_in_crop = int(np.median(foam_x_array[on_top_line]))
else:
top_x_in_crop = int(np.median(foam_x_array[top_region_mask]))
container_bottom = container_bottom_in_crop + crop_top
container_top = container_top_in_crop + crop_top
bottom_x = bottom_x_in_crop + crop_left
top_x = top_x_in_crop + crop_left
method = 'liquid_foam'
confidence = (liquid_conf + foam_conf) / 2
# 情况2: 仅liquid
elif has_liquid:
liquid_y_array = np.array(liquid_y_coords)
liquid_x_array = np.array(liquid_x_coords)
max_y = np.max(liquid_y_array)
bottom_region_mask = liquid_y_array >= (max_y - 5)
container_bottom_in_crop = int(np.median(liquid_y_array[bottom_region_mask]))
on_bottom_line = liquid_y_array == container_bottom_in_crop
if np.sum(on_bottom_line) > 0:
bottom_x_in_crop = int(np.median(liquid_x_array[on_bottom_line]))
else:
bottom_x_in_crop = int(np.median(liquid_x_array[bottom_region_mask]))
min_y = np.min(liquid_y_array)
top_region_mask = liquid_y_array <= (min_y + 5)
container_top_in_crop = int(np.median(liquid_y_array[top_region_mask]))
on_top_line = liquid_y_array == container_top_in_crop
if np.sum(on_top_line) > 0:
top_x_in_crop = int(np.median(liquid_x_array[on_top_line]))
else:
top_x_in_crop = int(np.median(liquid_x_array[top_region_mask]))
container_bottom = container_bottom_in_crop + crop_top
container_top = container_top_in_crop + crop_top
bottom_x = bottom_x_in_crop + crop_left
top_x = top_x_in_crop + crop_left
method = 'liquid_only'
confidence = liquid_conf
# 情况1: 仅air
elif has_air:
air_y_array = np.array(air_y_coords)
air_x_array = np.array(air_x_coords)
max_y = np.max(air_y_array)
bottom_region_mask = air_y_array >= (max_y - 5)
container_bottom_in_crop = int(np.median(air_y_array[bottom_region_mask]))
on_bottom_line = air_y_array == container_bottom_in_crop
if np.sum(on_bottom_line) > 0:
bottom_x_in_crop = int(np.median(air_x_array[on_bottom_line]))
else:
bottom_x_in_crop = int(np.median(air_x_array[bottom_region_mask]))
min_y = np.min(air_y_array)
top_region_mask = air_y_array <= (min_y + 5)
container_top_in_crop = int(np.median(air_y_array[top_region_mask]))
on_top_line = air_y_array == container_top_in_crop
if np.sum(on_top_line) > 0:
top_x_in_crop = int(np.median(air_x_array[on_top_line]))
else:
top_x_in_crop = int(np.median(air_x_array[top_region_mask]))
container_bottom = container_bottom_in_crop + crop_top
container_top = container_top_in_crop + crop_top
bottom_x = bottom_x_in_crop + crop_left
top_x = top_x_in_crop + crop_left
method = 'air_only'
confidence = air_conf
# 情况3: 仅foam
elif has_foam:
all_foam_y = []
all_foam_x = []
for mask, y_coords, x_coords in foam_masks:
all_foam_y.extend(y_coords)
all_foam_x.extend(x_coords)
foam_y_array = np.array(all_foam_y)
foam_x_array = np.array(all_foam_x)
max_y = np.max(foam_y_array)
bottom_region_mask = foam_y_array >= (max_y - 5)
container_bottom_in_crop = int(np.median(foam_y_array[bottom_region_mask]))
on_bottom_line = foam_y_array == container_bottom_in_crop
if np.sum(on_bottom_line) > 0:
bottom_x_in_crop = int(np.median(foam_x_array[on_bottom_line]))
else:
bottom_x_in_crop = int(np.median(foam_x_array[bottom_region_mask]))
min_y = np.min(foam_y_array)
top_region_mask = foam_y_array <= (min_y + 5)
container_top_in_crop = int(np.median(foam_y_array[top_region_mask]))
on_top_line = foam_y_array == container_top_in_crop
if np.sum(on_top_line) > 0:
top_x_in_crop = int(np.median(foam_x_array[on_top_line]))
else:
top_x_in_crop = int(np.median(foam_x_array[top_region_mask]))
container_bottom = container_bottom_in_crop + crop_top
container_top = container_top_in_crop + crop_top
bottom_x = bottom_x_in_crop + crop_left
top_x = top_x_in_crop + crop_left
method = 'foam_only'
confidence = foam_conf
else:
return None
if container_bottom is None or container_top is None:
return None
if container_bottom <= container_top:
return None
container_height = container_bottom - container_top
return {
'index': idx,
'top': int(container_top),
'bottom': int(container_bottom),
'top_x': int(top_x),
'bottom_x': int(bottom_x),
'height': int(container_height),
'confidence': float(confidence),
'method': method
}
def _draw_annotations(
self,
image: np.ndarray,
container_info: Dict,
left: int,
right: int
):
"""在图片上绘制标注点和线"""
top_y = container_info['top']
bottom_y = container_info['bottom']
top_x = container_info['top_x']
bottom_x = container_info['bottom_x']
idx = container_info['index']
top_color = (0, 255, 0)
bottom_color = (0, 0, 255)
line_color = (255, 255, 0)
cv2.circle(image, (top_x, top_y), 8, top_color, -1)
cv2.circle(image, (top_x, top_y), 10, top_color, 2)
cv2.circle(image, (bottom_x, bottom_y), 8, bottom_color, -1)
cv2.circle(image, (bottom_x, bottom_y), 10, bottom_color, 2)
cv2.line(image, (top_x, top_y), (bottom_x, bottom_y), line_color, 2)
cv2.line(image, (left, top_y), (right, top_y), top_color, 1, cv2.LINE_AA)
cv2.line(image, (left, bottom_y), (right, bottom_y), bottom_color, 1, cv2.LINE_AA)
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = 0.6
thickness = 2
cv2.putText(image, f"Top-{idx}", (top_x + 15, top_y - 10),
font, font_scale, top_color, thickness)
cv2.putText(image, f"Bottom-{idx}", (bottom_x + 15, bottom_y + 20),
font, font_scale, bottom_color, thickness)
mid_y = (top_y + bottom_y) // 2
mid_x = max(top_x, bottom_x) + 20
cv2.putText(image, f"{container_info['height']}px", (mid_x, mid_y),
font, font_scale, line_color, thickness)
def _draw_boxes(self, image: np.ndarray, boxes: List[Dict]):
"""在图片上绘制box"""
# 定义类别颜色
class_colors = {
'liquid': (0, 255, 0), # 绿色
'air': (255, 0, 0), # 蓝色
'foam': (0, 165, 255), # 橙色
'default': (255, 255, 0) # 青色
}
for box in boxes:
x1, y1, x2, y2 = box['x1'], box['y1'], box['x2'], box['y2']
classes = box['classes'] # 现在是列表
# 如果包含多个类别,使用默认颜色;否则使用类别颜色
if len(classes) > 1:
color = (0, 255, 255) # 黄色 - 表示合并框
label = f"Region ({'+'.join(classes)})"
else:
class_name = classes[0]
color = class_colors.get(class_name, class_colors['default'])
label = f"{class_name}"
# 绘制矩形框
cv2.rectangle(image, (x1, y1), (x2, y2), color, 3)
# 绘制标签背景
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = 0.7
thickness = 2
(text_width, text_height), baseline = cv2.getTextSize(
label, font, font_scale, thickness
)
# 标签背景矩形
cv2.rectangle(
image,
(x1, y1 - text_height - baseline - 8),
(x1 + text_width + 8, y1),
color,
-1
)
# 绘制标签文字
cv2.putText(
image,
label,
(x1 + 4, y1 - baseline - 4),
font,
font_scale,
(0, 0, 0), # 黑色文字
thickness
)
# 兼容性别名
AutoboxDetector = AutoAnnotationDetector
def test_auto_annotation():
"""测试自动标注功能(整合标框和标点)"""
import os
print("="*80)
print("🧪 自动标注功能测试")
print("="*80)
# 配置参数
model_path = r"D:\restructure\liquid_level_line_detection_system\database\model\detection_model\detect\best.dat"
test_image_path = r"D:\restructure\liquid_level_line_detection_system\test_image\4.jpg"
output_dir = r"D:\restructure\liquid_level_line_detection_system\test_output"
# 检查文件
if not os.path.exists(model_path):
print(f"❌ 模型文件不存在: {model_path}")
return
if not os.path.exists(test_image_path):
print(f"❌ 测试图片不存在: {test_image_path}")
print(f"💡 请将测试图片放置到: {test_image_path}")
return
# 创建输出目录
os.makedirs(output_dir, exist_ok=True)
# 加载图片
print(f"\n📷 加载测试图片: {test_image_path}")
image = cv2.imread(test_image_path)
if image is None:
print(f"❌ 无法读取图片")
return
print(f" 图片尺寸: {image.shape[1]}x{image.shape[0]}")
# 创建检测器
print(f"\n🔧 初始化自动标注检测器...")
detector = AutoAnnotationDetector(model_path=model_path, device='cuda')
# 执行核心检测
print(f"\n🚀 开始执行检测...")
detection_result = detector.detect(
image=image,
conf_threshold=0.5,
min_area=100
)
if not detection_result.get('success'):
print(f"❌ 检测失败")
if 'error' in detection_result:
print(f" 错误信息: {detection_result['error']}")
return
# ========== 生成box和点数据 ==========
print(f"\n{'='*80}")
print(f"📦 生成Box标注数据")
print(f"{'='*80}")
boxes = detector.get_boxes(
detection_result,
padding=10,
merge_all=True # True=合并所有类别, False=按类别分别生成
)
print(f"\n生成的Box区域:")
for i, box in enumerate(boxes, 1):
print(f" Box {i}:")
print(f" - 包含类别: {', '.join(box['classes'])}")
print(f" - 坐标: ({box['x1']}, {box['y1']}) -> ({box['x2']}, {box['y2']})")
print(f" - 面积: {box['area']}px²")
print(f"\n{'='*80}")
print(f"📍 生成Point标注数据")
print(f"{'='*80}")
points = detector.get_points(detection_result)
print(f"\n生成的Point位置:")
for i, point in enumerate(points, 1):
print(f" Point {i}:")
print(f" - 顶部: ({point['top_x']}, {point['top']})")
print(f" - 底部: ({point['bottom_x']}, {point['bottom']})")
print(f" - 高度: {point['height']}px")
print(f" - 检测方法: {point['method']}")
# ========== 绘制并保存结果 ==========
print(f"\n{'='*80}")
print(f"🎨 绘制标注结果")
print(f"{'='*80}")
# 在同一张图上绘制框和点
combined_image = detector.draw_all(image, boxes, points)
combined_output_path = os.path.join(output_dir, "auto_annotation_result.jpg")
cv2.imwrite(combined_output_path, combined_image)
print(f"\n💾 完整标注图片已保存: {combined_output_path}")
# 显示图片(可选)
try:
cv2.imshow("Auto Annotation Result (Box + Point)", combined_image)
print(f"\n👀 按任意键关闭图片窗口...")
cv2.waitKey(0)
cv2.destroyAllWindows()
except:
print(f"⚠️ 无法显示图片(可能是无GUI环境)")
print(f"\n{'='*80}")
print(f"✅ 测试完成")
print(f"{'='*80}")
if __name__ == "__main__": if __name__ == "__main__":
test_auto_dot() # 测试自动标点功能(原始版本)
# test_auto_dot()
# 测试整合的自动标注功能(标框+标点)
test_auto_annotation()
...@@ -118,7 +118,7 @@ class TrainingWorker(QThread): ...@@ -118,7 +118,7 @@ class TrainingWorker(QThread):
# 如果签名不匹配,说明这是一个直接重命名的 .pt 文件 # 如果签名不匹配,说明这是一个直接重命名的 .pt 文件
if signature != MODEL_FILE_SIGNATURE: if signature != MODEL_FILE_SIGNATURE:
print(f"[警告] {dat_path.name} 不是加密的 .dat 文件,将直接作为 .pt 文件使用") print(f"[警告] {dat_path.name} 不是加密的 .dat 文件,将直接作为 .pt 文件使用")
# 直接返回原路径,YOLO 可以直接加载 # 直接返回原路径,模型 可以直接加载
return str(dat_path) return str(dat_path)
# 继续解密流程 # 继续解密流程
......
1曲线模式索引0布局,只显示根据curvemission筛选使用的通道面板失效了 1曲线模式索引0布局,只显示根据curvemission筛选使用的通道面板失效了
...@@ -25,4 +25,38 @@ ...@@ -25,4 +25,38 @@
6mission_status判断逻辑,只要任意channelmission值为此任务,则此任务mission_status=true 6mission_status判断逻辑,只要任意channelmission值为此任务,则此任务mission_status=true
class ModelLoadingProgressDialog(QDialog): class ModelLoadingProgressDialog(QDialog):
"""模型加载进度条对话框""" """模型加载进度条对话框"""
\ No newline at end of file 自动选择标注区域类
1.设置一个或多个最小面积box,box区域包含了某一区域的分割结果
自动标点,根据分割结果选择对应分支逻辑
1: liquid_only
- **容器底部**: liquid掩码的最低点
- **容器顶部**: liquid掩码的最高点
- **适用场景**: 只检测到液体,未检测到空气和泡沫
2: air_only
- **容器底部**: air掩码的最低点
- **容器顶部**: air掩码的最高点
- **适用场景**: 只检测到空气,未检测到液体和泡沫
3: foam_only
- **容器底部**: foam掩码的最低点
- **容器顶部**: foam掩码的最高点
- **适用场景**: 只检测到泡沫,未检测到液体和空气
4: liquid_air
- **容器底部**: liquid掩码的最低点
- **容器顶部**: air掩码的最高点
- **适用场景**: 同时检测到液体和空气,未检测到泡沫
5: liquid_foam
- **容器底部**: liquid掩码的最低点
- **容器顶部**: foam掩码的最高点
- **适用场景**: 同时检测到液体和泡沫,未检测到空气
6: liquid_foam_air
- **容器底部**: liquid掩码的最低点
- **容器顶部**: air掩码的最高点
- **适用场景**: 同时检测到液体、泡沫和空气
7: foam_air
- **容器底部**: foam掩码的最低点
- **容器顶部**: air掩码的最高点
- **适用场景**: 同时检测到泡沫和空气,未检测到液体
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment