提高视频帧率
This commit is contained in:
@@ -12,7 +12,7 @@ from app.models import Detection
|
||||
class PythonModelDetector:
|
||||
"""Object detector using native Python models"""
|
||||
|
||||
def __init__(self, model_name: str, model_path: str, input_width: int, input_height: int, color: int = 0x00FF00):
|
||||
def __init__(self, model_name: str, model_path: str, input_width: int, input_height: int, color: int = 0x00FF00, model_config: dict = None):
|
||||
"""
|
||||
Initialize detector with Python model
|
||||
|
||||
@@ -22,11 +22,13 @@ class PythonModelDetector:
|
||||
input_width: Input width for the model
|
||||
input_height: Input height for the model
|
||||
color: RGB color for detection boxes (default: green)
|
||||
model_config: Additional configuration to pass to the model
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self.input_width = input_width
|
||||
self.input_height = input_height
|
||||
self.color = color
|
||||
self.model_config = model_config or {}
|
||||
|
||||
# Convert color from RGB to BGR (OpenCV uses BGR)
|
||||
self.color_bgr = ((color & 0xFF) << 16) | (color & 0xFF00) | ((color >> 16) & 0xFF)
|
||||
@@ -72,8 +74,18 @@ class PythonModelDetector:
|
||||
if not hasattr(model_module, "Model"):
|
||||
raise AttributeError(f"Model module must define a 'Model' class: {model_path}")
|
||||
|
||||
# Create model instance
|
||||
self.model = model_module.Model()
|
||||
# Create model instance with config
|
||||
# Try to pass config to model constructor if it accepts parameters
|
||||
import inspect
|
||||
model_class = model_module.Model
|
||||
sig = inspect.signature(model_class.__init__)
|
||||
|
||||
if len(sig.parameters) > 1: # Has parameters beyond 'self'
|
||||
# Pass all config as keyword arguments
|
||||
self.model = model_class(**self.model_config)
|
||||
else:
|
||||
# No parameters, create without arguments
|
||||
self.model = model_class()
|
||||
|
||||
# Check if model has the required methods
|
||||
if not hasattr(self.model, "predict"):
|
||||
@@ -113,18 +125,16 @@ class PythonModelDetector:
|
||||
# Original image dimensions
|
||||
img_height, img_width = img.shape[:2]
|
||||
|
||||
# Preprocess image
|
||||
processed_img = self.preprocess(img)
|
||||
|
||||
# Measure inference time
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Run inference using model's predict method
|
||||
# Note: Pass original image to model, let it handle preprocessing
|
||||
# Expected return format from model's predict:
|
||||
# List of dicts with keys: 'bbox', 'class_id', 'confidence'
|
||||
# bbox: (x, y, w, h) normalized [0-1]
|
||||
model_results = self.model.predict(processed_img)
|
||||
model_results = self.model.predict(img)
|
||||
|
||||
# Calculate inference time in milliseconds
|
||||
inference_time = (time.time() - start_time) * 1000
|
||||
@@ -279,13 +289,22 @@ class ModelManager:
|
||||
# Use color from palette
|
||||
color = palette[i % len(palette)]
|
||||
|
||||
# Extract model-specific config (model_file, model_name, etc.)
|
||||
# These will be passed to the Model class __init__
|
||||
model_init_config = {}
|
||||
if "model_file" in model_config:
|
||||
model_init_config["model_file"] = model_config["model_file"]
|
||||
if "display_name" in model_config:
|
||||
model_init_config["model_name"] = model_config["display_name"]
|
||||
|
||||
# Create detector for Python model
|
||||
detector = PythonModelDetector(
|
||||
model_name=name,
|
||||
model_path=path,
|
||||
input_width=size[0],
|
||||
input_height=size[1],
|
||||
color=color
|
||||
color=color,
|
||||
model_config=model_init_config
|
||||
)
|
||||
|
||||
self.models[name] = detector
|
||||
|
||||
@@ -129,17 +129,13 @@ async def detect_file(
|
||||
model_name: str,
|
||||
file: UploadFile = File(...)
|
||||
):
|
||||
# 1. 打印 model_name（直接打印字符串即可）
|
||||
print(f"接收到的 model_name: {model_name}")
|
||||
|
||||
# 2. 打印 file 的基本信息（文件名、内容类型等）
|
||||
print(f"文件名: {file.filename}")
|
||||
print(f"文件内容类型: {file.content_type}") # 例如 image/jpeg、text/plain 等
|
||||
"""Detect objects in an uploaded image file"""
|
||||
print(f"接收到的 model_name: {model_name}")
|
||||
print(f"文件名: {file.filename}")
|
||||
print(f"文件内容类型: {file.content_type}")
|
||||
|
||||
global model_manager
|
||||
|
||||
|
||||
|
||||
if not model_manager:
|
||||
raise HTTPException(status_code=500, detail="Model manager not initialized")
|
||||
|
||||
@@ -151,16 +147,39 @@ async def detect_file(
|
||||
# Read uploaded file
|
||||
try:
|
||||
contents = await file.read()
|
||||
print(f"文件大小: {len(contents)} 字节")
|
||||
|
||||
if len(contents) == 0:
|
||||
raise HTTPException(status_code=400, detail="Empty file")
|
||||
|
||||
nparr = np.frombuffer(contents, np.uint8)
|
||||
print(f"numpy数组形状: {nparr.shape}, dtype: {nparr.dtype}")
|
||||
|
||||
image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
||||
|
||||
if image is None:
|
||||
raise HTTPException(status_code=400, detail="Invalid image data")
|
||||
print("错误: cv2.imdecode 返回 None")
|
||||
raise HTTPException(status_code=400, detail="Invalid image data - failed to decode")
|
||||
|
||||
print(f"解码后图像形状: {image.shape}, dtype: {image.dtype}")
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"处理图像时出错: {str(e)}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
raise HTTPException(status_code=400, detail=f"Failed to process image: {str(e)}")
|
||||
|
||||
# Run detection
|
||||
try:
|
||||
detections, inference_time = detector.detect(image)
|
||||
print(f"检测完成: 找到 {len(detections)} 个目标, 耗时 {inference_time:.2f}ms")
|
||||
except Exception as e:
|
||||
print(f"推理过程中出错: {str(e)}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=f"Detection failed: {str(e)}")
|
||||
|
||||
return DetectionResponse(
|
||||
model_name=model_name,
|
||||
|
||||
@@ -1,14 +1,18 @@
|
||||
[
|
||||
{
|
||||
"name": "smoke",
|
||||
"path": "models/smoke_model.py",
|
||||
"path": "models/universal_yolo_model.py",
|
||||
"model_file": "smoke.pt",
|
||||
"display_name": "吸烟检测",
|
||||
"size": [640, 640],
|
||||
"comment": "烟雾检测模型"
|
||||
"comment": "吸烟检测模型 - YOLOv11"
|
||||
},
|
||||
{
|
||||
"name": "garbage",
|
||||
"path": "models/garbage_model.py",
|
||||
"path": "models/universal_yolo_model.py",
|
||||
"model_file": "garbage.pt",
|
||||
"display_name": "垃圾识别",
|
||||
"size": [640, 640],
|
||||
"comment": "垃圾检测模型"
|
||||
"comment": "垃圾检测模型 - YOLOv8"
|
||||
}
|
||||
]
|
||||
@@ -1,211 +0,0 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import cv2
|
||||
from typing import List, Dict, Any
|
||||
import torch
|
||||
|
||||
class Model:
|
||||
"""
|
||||
垃圾识别模型 - 直接加载 PyTorch 模型文件
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""初始化模型"""
|
||||
# 获取当前文件所在目录路径
|
||||
model_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
# 模型文件路径
|
||||
model_path = os.path.join(model_dir, "best.pt")
|
||||
|
||||
print(f"正在加载垃圾识别模型: {model_path}")
|
||||
|
||||
# 加载 PyTorch 模型
|
||||
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"使用设备: {self.device}")
|
||||
|
||||
# 使用 YOLOv5 或通用方式加载模型
|
||||
try:
|
||||
# 尝试使用 YOLOv5 加载
|
||||
import sys
|
||||
sys.path.append(os.path.dirname(model_dir)) # 添加父目录到路径
|
||||
|
||||
try:
|
||||
# 方法1: 如果安装了 YOLOv5
|
||||
import yolov5
|
||||
self.model = yolov5.load(model_path, device=self.device)
|
||||
self.yolov5_api = True
|
||||
print("使用 YOLOv5 包加载模型")
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
# 方法2: 直接加载 YOLO 代码
|
||||
from models.yolov5_utils import attempt_load
|
||||
self.model = attempt_load(model_path, device=self.device)
|
||||
self.yolov5_api = False
|
||||
print("使用内置 YOLOv5 工具加载模型")
|
||||
|
||||
except Exception as e:
|
||||
# 方法3: 通用 PyTorch 加载
|
||||
print(f"YOLOv5 加载失败: {e}")
|
||||
print("使用通用 PyTorch 加载")
|
||||
self.model = torch.load(
|
||||
model_path,
|
||||
map_location=self.device,
|
||||
weights_only=False # 允许加载模型类结构（解决 PyTorch 2.6+ 兼容性问题）
|
||||
)
|
||||
if isinstance(self.model, dict) and 'model' in self.model:
|
||||
self.model = self.model['model']
|
||||
self.yolov5_api = False
|
||||
|
||||
# 如果是 ScriptModule，设置为评估模式
|
||||
if isinstance(self.model, torch.jit.ScriptModule):
|
||||
self.model.eval()
|
||||
elif hasattr(self.model, 'eval'):
|
||||
self.model.eval()
|
||||
|
||||
# 加载类别名称
|
||||
self.classes = []
|
||||
classes_path = os.path.join(model_dir, "classes.txt")
|
||||
if os.path.exists(classes_path):
|
||||
with open(classes_path, 'r', encoding='utf-8') as f:
|
||||
self.classes = [line.strip() for line in f.readlines() if line.strip()]
|
||||
print(f"已加载 {len(self.classes)} 个类别")
|
||||
else:
|
||||
# 如果模型自带类别信息
|
||||
if hasattr(self.model, 'names') and self.model.names:
|
||||
self.classes = self.model.names
|
||||
print(f"使用模型自带类别，共 {len(self.classes)} 个类别")
|
||||
else:
|
||||
print("未找到类别文件，将使用数字索引作为类别名")
|
||||
|
||||
# 设置识别参数
|
||||
self.conf_threshold = 0.25 # 置信度阈值
|
||||
self.img_size = 640 # 默认输入图像大小
|
||||
|
||||
print("垃圾识别模型加载完成")
|
||||
|
||||
def preprocess(self, image: np.ndarray) -> np.ndarray:
|
||||
"""预处理图像"""
|
||||
# 如果是使用 YOLOv5 API，不需要预处理
|
||||
if hasattr(self, 'yolov5_api') and self.yolov5_api:
|
||||
return image
|
||||
|
||||
# 默认预处理：调整大小并归一化
|
||||
img = cv2.resize(image, (self.img_size, self.img_size))
|
||||
|
||||
# BGR 转 RGB
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
|
||||
# 归一化 [0, 255] -> [0, 1]
|
||||
img = img / 255.0
|
||||
|
||||
# HWC -> CHW (高度，宽度，通道 -> 通道，高度，宽度)
|
||||
img = img.transpose(2, 0, 1)
|
||||
|
||||
# 转为 torch tensor
|
||||
img = torch.from_numpy(img).float()
|
||||
|
||||
# 添加批次维度
|
||||
img = img.unsqueeze(0)
|
||||
|
||||
# 移至设备
|
||||
img = img.to(self.device)
|
||||
|
||||
return img
|
||||
|
||||
def predict(self, image: np.ndarray) -> List[Dict[str, Any]]:
|
||||
"""模型推理"""
|
||||
original_height, original_width = image.shape[:2]
|
||||
|
||||
try:
|
||||
# 如果使用 YOLOv5 API
|
||||
if hasattr(self, 'yolov5_api') and self.yolov5_api:
|
||||
# YOLOv5 API 直接处理图像
|
||||
results = self.model(image)
|
||||
|
||||
# 提取检测结果
|
||||
predictions = results.pred[0] # 第一批次的预测
|
||||
|
||||
detections = []
|
||||
for *xyxy, conf, cls_id in predictions.cpu().numpy():
|
||||
x1, y1, x2, y2 = xyxy
|
||||
|
||||
# 转换为归一化坐标 (x, y, w, h)
|
||||
x = x1 / original_width
|
||||
y = y1 / original_height
|
||||
w = (x2 - x1) / original_width
|
||||
h = (y2 - y1) / original_height
|
||||
|
||||
# 整数类别 ID
|
||||
cls_id = int(cls_id)
|
||||
|
||||
# 获取类别名称
|
||||
class_name = f"cls{cls_id}"
|
||||
if 0 <= cls_id < len(self.classes):
|
||||
class_name = self.classes[cls_id]
|
||||
|
||||
# 添加检测结果
|
||||
if conf >= self.conf_threshold:
|
||||
detections.append({
|
||||
'bbox': (x, y, w, h),
|
||||
'class_id': cls_id,
|
||||
'confidence': float(conf)
|
||||
})
|
||||
|
||||
return detections
|
||||
|
||||
else:
|
||||
# 通用 PyTorch 模型处理
|
||||
# 预处理图像
|
||||
img = self.preprocess(image)
|
||||
|
||||
# 推理
|
||||
with torch.no_grad():
|
||||
outputs = self.model(img)
|
||||
|
||||
# 后处理结果（这里需要根据模型输出格式调整）
|
||||
detections = []
|
||||
|
||||
# 假设输出格式是 YOLO 风格：[batch_idx, x1, y1, x2, y2, conf, cls_id]
|
||||
if isinstance(outputs, torch.Tensor) and outputs.dim() == 2 and outputs.size(1) >= 6:
|
||||
for *xyxy, conf, cls_id in outputs.cpu().numpy():
|
||||
if conf >= self.conf_threshold:
|
||||
x1, y1, x2, y2 = xyxy
|
||||
|
||||
# 转换为归一化坐标 (x, y, w, h)
|
||||
x = x1 / original_width
|
||||
y = y1 / original_height
|
||||
w = (x2 - x1) / original_width
|
||||
h = (y2 - y1) / original_height
|
||||
|
||||
# 整数类别 ID
|
||||
cls_id = int(cls_id)
|
||||
|
||||
detections.append({
|
||||
'bbox': (x, y, w, h),
|
||||
'class_id': cls_id,
|
||||
'confidence': float(conf)
|
||||
})
|
||||
# 处理其他可能的输出格式
|
||||
else:
|
||||
# 这里需要根据模型的实际输出格式进行适配
|
||||
print("警告：无法识别的模型输出格式，请检查模型类型")
|
||||
|
||||
return detections
|
||||
|
||||
except Exception as e:
|
||||
print(f"推理过程中出错: {str(e)}")
|
||||
# 出错时返回空结果
|
||||
return []
|
||||
|
||||
@property
|
||||
def applies_nms(self) -> bool:
|
||||
"""模型是否内部应用了 NMS"""
|
||||
# YOLOv5 会自动应用 NMS
|
||||
return True
|
||||
|
||||
def close(self):
|
||||
"""释放资源"""
|
||||
if hasattr(self, 'model'):
|
||||
# 删除模型以释放 GPU 内存
|
||||
del self.model
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
print("垃圾识别模型已关闭")
|
||||
1
python-inference-service/models/smoke.txt
Normal file
1
python-inference-service/models/smoke.txt
Normal file
@@ -0,0 +1 @@
|
||||
垃圾
|
||||
@@ -1,126 +0,0 @@
|
||||
import numpy as np
|
||||
import cv2
|
||||
from typing import List, Dict, Any, Tuple
|
||||
|
||||
class Model:
|
||||
"""
|
||||
Smoke detection model implementation
|
||||
|
||||
This is a simple example that could be replaced with an actual
|
||||
TensorFlow, PyTorch, or other ML framework implementation.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize smoke detection model"""
|
||||
# In a real implementation, you would load your model here
|
||||
print("Smoke detection model initialized")
|
||||
|
||||
# Define smoke class IDs
|
||||
self.smoke_classes = {
|
||||
0: "smoke",
|
||||
1: "fire"
|
||||
}
|
||||
|
||||
def preprocess(self, image: np.ndarray) -> np.ndarray:
|
||||
"""Preprocess image for model input"""
|
||||
# Convert BGR to grayscale for smoke detection
|
||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||||
# Convert back to 3 channels to match model expected input shape
|
||||
gray_3ch = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
# In a real implementation, you would do normalization, etc.
|
||||
return gray_3ch
|
||||
|
||||
def predict(self, image: np.ndarray) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Run smoke detection on the image
|
||||
|
||||
This is a simplified example that uses basic image processing
|
||||
In a real implementation, you would use your ML model
|
||||
"""
|
||||
# Convert to grayscale for processing
|
||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
# Apply Gaussian blur to reduce noise
|
||||
blurred = cv2.GaussianBlur(gray, (15, 15), 0)
|
||||
|
||||
# Simple thresholding to find potential smoke regions
|
||||
# In a real implementation, you'd use a trained model
|
||||
_, thresh = cv2.threshold(blurred, 100, 255, cv2.THRESH_BINARY)
|
||||
|
||||
# Find contours in the thresholded image
|
||||
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
# Process contours to find potential smoke regions
|
||||
detections = []
|
||||
height, width = image.shape[:2]
|
||||
|
||||
for contour in contours:
|
||||
# Get bounding box
|
||||
x, y, w, h = cv2.boundingRect(contour)
|
||||
|
||||
# Filter small regions
|
||||
if w > width * 0.05 and h > height * 0.05:
|
||||
# Calculate area ratio
|
||||
area = cv2.contourArea(contour)
|
||||
rect_area = w * h
|
||||
fill_ratio = area / rect_area if rect_area > 0 else 0
|
||||
|
||||
# Smoke tends to have irregular shapes
|
||||
# This is just for demonstration purposes
|
||||
if fill_ratio > 0.2 and fill_ratio < 0.8:
|
||||
# Normalize coordinates
|
||||
x_norm = x / width
|
||||
y_norm = y / height
|
||||
w_norm = w / width
|
||||
h_norm = h / height
|
||||
|
||||
# Determine if it's smoke or fire (just a simple heuristic for demo)
|
||||
# In a real model, this would be determined by the model prediction
|
||||
class_id = 0 # Default to smoke
|
||||
|
||||
# Check if the region has high red values (fire)
|
||||
roi = image[y:y+h, x:x+w]
|
||||
if roi.size > 0: # Make sure ROI is not empty
|
||||
avg_color = np.mean(roi, axis=(0, 1))
|
||||
if avg_color[2] > 150 and avg_color[2] > avg_color[0] * 1.5: # High red, indicating fire
|
||||
class_id = 1 # Fire
|
||||
|
||||
# Calculate confidence based on fill ratio
|
||||
# This is just for demonstration
|
||||
confidence = 0.5 + fill_ratio * 0.3
|
||||
|
||||
# Add to detections
|
||||
detections.append({
|
||||
'bbox': (x_norm, y_norm, w_norm, h_norm),
|
||||
'class_id': class_id,
|
||||
'confidence': confidence
|
||||
})
|
||||
|
||||
# For demo purposes, if no smoke detected by algorithm,
|
||||
# add a small chance of random detection
|
||||
if not detections and np.random.random() < 0.1: # 10% chance
|
||||
# Random smoke detection
|
||||
x = np.random.random() * 0.7
|
||||
y = np.random.random() * 0.7
|
||||
w = 0.1 + np.random.random() * 0.2
|
||||
h = 0.1 + np.random.random() * 0.2
|
||||
confidence = 0.5 + np.random.random() * 0.3
|
||||
|
||||
detections.append({
|
||||
'bbox': (x, y, w, h),
|
||||
'class_id': 0, # Smoke
|
||||
'confidence': confidence
|
||||
})
|
||||
|
||||
return detections
|
||||
|
||||
@property
|
||||
def applies_nms(self) -> bool:
|
||||
"""Model does not apply NMS internally"""
|
||||
return False
|
||||
|
||||
def close(self):
|
||||
"""Release resources"""
|
||||
# In a real implementation, you would release model resources here
|
||||
pass
|
||||
@@ -1,211 +0,0 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import cv2
|
||||
from typing import List, Dict, Any
|
||||
import torch
|
||||
|
||||
class Model:
|
||||
"""
|
||||
垃圾识别模型 - 直接加载 PyTorch 模型文件
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""初始化模型"""
|
||||
# 获取当前文件所在目录路径
|
||||
model_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
# 模型文件路径
|
||||
model_path = os.path.join(model_dir, "smoke.pt")
|
||||
|
||||
print(f"正在加载垃圾识别模型: {model_path}")
|
||||
|
||||
# 加载 PyTorch 模型
|
||||
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"使用设备: {self.device}")
|
||||
|
||||
# 使用 YOLOv5 或通用方式加载模型
|
||||
try:
|
||||
# 尝试使用 YOLOv5 加载
|
||||
import sys
|
||||
sys.path.append(os.path.dirname(model_dir)) # 添加父目录到路径
|
||||
|
||||
try:
|
||||
# 方法1: 如果安装了 YOLOv5
|
||||
import yolov5
|
||||
self.model = yolov5.load(model_path, device=self.device)
|
||||
self.yolov5_api = True
|
||||
print("使用 YOLOv5 包加载模型")
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
# 方法2: 直接加载 YOLO 代码
|
||||
from models.yolov5_utils import attempt_load
|
||||
self.model = attempt_load(model_path, device=self.device)
|
||||
self.yolov5_api = False
|
||||
print("使用内置 YOLOv5 工具加载模型")
|
||||
|
||||
except Exception as e:
|
||||
# 方法3: 通用 PyTorch 加载
|
||||
print(f"YOLOv5 加载失败: {e}")
|
||||
print("使用通用 PyTorch 加载")
|
||||
self.model = torch.load(
|
||||
model_path,
|
||||
map_location=self.device,
|
||||
weights_only=False # 允许加载模型类结构，解决 PyTorch 2.6+ 兼容性问题
|
||||
)
|
||||
if isinstance(self.model, dict) and 'model' in self.model:
|
||||
self.model = self.model['model']
|
||||
self.yolov5_api = False
|
||||
|
||||
# 如果是 ScriptModule，设置为评估模式
|
||||
if isinstance(self.model, torch.jit.ScriptModule):
|
||||
self.model.eval()
|
||||
elif hasattr(self.model, 'eval'):
|
||||
self.model.eval()
|
||||
|
||||
# 加载类别名称
|
||||
self.classes = []
|
||||
classes_path = os.path.join(model_dir, "classes.txt")
|
||||
if os.path.exists(classes_path):
|
||||
with open(classes_path, 'r', encoding='utf-8') as f:
|
||||
self.classes = [line.strip() for line in f.readlines() if line.strip()]
|
||||
print(f"已加载 {len(self.classes)} 个类别")
|
||||
else:
|
||||
# 如果模型自带类别信息
|
||||
if hasattr(self.model, 'names') and self.model.names:
|
||||
self.classes = self.model.names
|
||||
print(f"使用模型自带类别，共 {len(self.classes)} 个类别")
|
||||
else:
|
||||
print("未找到类别文件，将使用数字索引作为类别名")
|
||||
|
||||
# 设置识别参数
|
||||
self.conf_threshold = 0.25 # 置信度阈值
|
||||
self.img_size = 640 # 默认输入图像大小
|
||||
|
||||
print("垃圾识别模型加载完成")
|
||||
|
||||
def preprocess(self, image: np.ndarray) -> np.ndarray:
|
||||
"""预处理图像"""
|
||||
# 如果是使用 YOLOv5 API，不需要预处理
|
||||
if hasattr(self, 'yolov5_api') and self.yolov5_api:
|
||||
return image
|
||||
|
||||
# 默认预处理：调整大小并归一化
|
||||
img = cv2.resize(image, (self.img_size, self.img_size))
|
||||
|
||||
# BGR 转 RGB
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
|
||||
# 归一化 [0, 255] -> [0, 1]
|
||||
img = img / 255.0
|
||||
|
||||
# HWC -> CHW (高度，宽度，通道 -> 通道，高度，宽度)
|
||||
img = img.transpose(2, 0, 1)
|
||||
|
||||
# 转为 torch tensor
|
||||
img = torch.from_numpy(img).float()
|
||||
|
||||
# 添加批次维度
|
||||
img = img.unsqueeze(0)
|
||||
|
||||
# 移至设备
|
||||
img = img.to(self.device)
|
||||
|
||||
return img
|
||||
|
||||
def predict(self, image: np.ndarray) -> List[Dict[str, Any]]:
|
||||
"""模型推理"""
|
||||
original_height, original_width = image.shape[:2]
|
||||
|
||||
try:
|
||||
# 如果使用 YOLOv5 API
|
||||
if hasattr(self, 'yolov5_api') and self.yolov5_api:
|
||||
# YOLOv5 API 直接处理图像
|
||||
results = self.model(image)
|
||||
|
||||
# 提取检测结果
|
||||
predictions = results.pred[0] # 第一批次的预测
|
||||
|
||||
detections = []
|
||||
for *xyxy, conf, cls_id in predictions.cpu().numpy():
|
||||
x1, y1, x2, y2 = xyxy
|
||||
|
||||
# 转换为归一化坐标 (x, y, w, h)
|
||||
x = x1 / original_width
|
||||
y = y1 / original_height
|
||||
w = (x2 - x1) / original_width
|
||||
h = (y2 - y1) / original_height
|
||||
|
||||
# 整数类别 ID
|
||||
cls_id = int(cls_id)
|
||||
|
||||
# 获取类别名称
|
||||
class_name = f"cls{cls_id}"
|
||||
if 0 <= cls_id < len(self.classes):
|
||||
class_name = self.classes[cls_id]
|
||||
|
||||
# 添加检测结果
|
||||
if conf >= self.conf_threshold:
|
||||
detections.append({
|
||||
'bbox': (x, y, w, h),
|
||||
'class_id': cls_id,
|
||||
'confidence': float(conf)
|
||||
})
|
||||
|
||||
return detections
|
||||
|
||||
else:
|
||||
# 通用 PyTorch 模型处理
|
||||
# 预处理图像
|
||||
img = self.preprocess(image)
|
||||
|
||||
# 推理
|
||||
with torch.no_grad():
|
||||
outputs = self.model(img)
|
||||
|
||||
# 后处理结果（这里需要根据模型输出格式调整）
|
||||
detections = []
|
||||
|
||||
# 假设输出格式是 YOLO 风格：[batch_idx, x1, y1, x2, y2, conf, cls_id]
|
||||
if isinstance(outputs, torch.Tensor) and outputs.dim() == 2 and outputs.size(1) >= 6:
|
||||
for *xyxy, conf, cls_id in outputs.cpu().numpy():
|
||||
if conf >= self.conf_threshold:
|
||||
x1, y1, x2, y2 = xyxy
|
||||
|
||||
# 转换为归一化坐标 (x, y, w, h)
|
||||
x = x1 / original_width
|
||||
y = y1 / original_height
|
||||
w = (x2 - x1) / original_width
|
||||
h = (y2 - y1) / original_height
|
||||
|
||||
# 整数类别 ID
|
||||
cls_id = int(cls_id)
|
||||
|
||||
detections.append({
|
||||
'bbox': (x, y, w, h),
|
||||
'class_id': cls_id,
|
||||
'confidence': float(conf)
|
||||
})
|
||||
# 处理其他可能的输出格式
|
||||
else:
|
||||
# 这里需要根据模型的实际输出格式进行适配
|
||||
print("警告：无法识别的模型输出格式，请检查模型类型")
|
||||
|
||||
return detections
|
||||
|
||||
except Exception as e:
|
||||
print(f"推理过程中出错: {str(e)}")
|
||||
# 出错时返回空结果
|
||||
return []
|
||||
|
||||
@property
|
||||
def applies_nms(self) -> bool:
|
||||
"""模型是否内部应用了 NMS"""
|
||||
# YOLOv5 会自动应用 NMS
|
||||
return True
|
||||
|
||||
def close(self):
|
||||
"""释放资源"""
|
||||
if hasattr(self, 'model'):
|
||||
# 删除模型以释放 GPU 内存
|
||||
del self.model
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
print("垃圾识别模型已关闭")
|
||||
@@ -6,17 +6,39 @@ import torch
|
||||
|
||||
class Model:
|
||||
"""
|
||||
YOLOv8 模型包装类 - 使用 Ultralytics YOLO
|
||||
通用 YOLO 模型 - 支持 YOLOv8/YOLOv11 等基于 Ultralytics 的模型
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""初始化YOLOv8模型"""
|
||||
def __init__(self, model_file: str = None, model_name: str = "YOLO"):
|
||||
"""
|
||||
初始化模型
|
||||
|
||||
Args:
|
||||
model_file: 模型文件名（如 smoke.pt, best.pt）
|
||||
model_name: 模型显示名称（用于日志）
|
||||
"""
|
||||
# 获取当前文件所在目录路径
|
||||
model_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
# 模型文件路径
|
||||
model_path = os.path.join(model_dir, "best.pt")
|
||||
|
||||
print(f"正在加载YOLOv8模型: {model_path}")
|
||||
# 如果没有指定模型文件，尝试常见的文件名
|
||||
if model_file is None:
|
||||
for possible_file in ['garbage.pt', 'smoke.pt', 'best.pt', 'yolov8.pt', 'model.pt']:
|
||||
test_path = os.path.join(model_dir, possible_file)
|
||||
if os.path.exists(test_path):
|
||||
model_file = possible_file
|
||||
break
|
||||
|
||||
if model_file is None:
|
||||
raise FileNotFoundError(f"未找到模型文件，请在初始化时指定 model_file 参数")
|
||||
|
||||
# 模型文件路径
|
||||
model_path = os.path.join(model_dir, model_file)
|
||||
|
||||
if not os.path.exists(model_path):
|
||||
raise FileNotFoundError(f"模型文件不存在: {model_path}")
|
||||
|
||||
self.model_name = model_name
|
||||
print(f"正在加载{model_name}模型: {model_path}")
|
||||
|
||||
# 检查设备
|
||||
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
@@ -26,19 +48,30 @@ class Model:
|
||||
try:
|
||||
from ultralytics import YOLO
|
||||
self.model = YOLO(model_path)
|
||||
print("使用 Ultralytics YOLO 加载模型成功")
|
||||
print(f"使用 Ultralytics YOLO 加载模型成功")
|
||||
except ImportError:
|
||||
raise ImportError("请安装 ultralytics: pip install ultralytics>=8.0.0")
|
||||
except Exception as e:
|
||||
raise Exception(f"加载YOLOv8模型失败: {str(e)}")
|
||||
raise Exception(f"加载{model_name}模型失败: {str(e)}")
|
||||
|
||||
# 加载类别名称
|
||||
self.classes = []
|
||||
classes_path = os.path.join(model_dir, "classes.txt")
|
||||
if os.path.exists(classes_path):
|
||||
with open(classes_path, 'r', encoding='utf-8') as f:
|
||||
|
||||
# 1. 首先尝试加载与模型文件同名的类别文件（如 smoke.txt）
|
||||
model_base_name = os.path.splitext(model_file)[0]
|
||||
classes_path_specific = os.path.join(model_dir, f"{model_base_name}.txt")
|
||||
|
||||
# 2. 然后尝试加载通用的 classes.txt
|
||||
classes_path_generic = os.path.join(model_dir, "classes.txt")
|
||||
|
||||
if os.path.exists(classes_path_specific):
|
||||
with open(classes_path_specific, 'r', encoding='utf-8') as f:
|
||||
self.classes = [line.strip() for line in f.readlines() if line.strip()]
|
||||
print(f"已加载 {len(self.classes)} 个类别")
|
||||
print(f"已加载类别文件: {model_base_name}.txt ({len(self.classes)} 个类别)")
|
||||
elif os.path.exists(classes_path_generic):
|
||||
with open(classes_path_generic, 'r', encoding='utf-8') as f:
|
||||
self.classes = [line.strip() for line in f.readlines() if line.strip()]
|
||||
print(f"已加载类别文件: classes.txt ({len(self.classes)} 个类别)")
|
||||
else:
|
||||
# 使用模型自带的类别信息
|
||||
if hasattr(self.model, 'names') and self.model.names:
|
||||
@@ -51,10 +84,10 @@ class Model:
|
||||
self.conf_threshold = 0.25 # 置信度阈值
|
||||
self.img_size = 640 # 默认输入图像大小
|
||||
|
||||
print("YOLOv8模型加载完成")
|
||||
print(f"{model_name}模型加载完成")
|
||||
|
||||
def preprocess(self, image: np.ndarray) -> np.ndarray:
|
||||
"""预处理图像 - YOLOv8会自动处理，这里直接返回"""
|
||||
"""预处理图像 - Ultralytics YOLO 会自动处理，这里直接返回"""
|
||||
return image
|
||||
|
||||
def predict(self, image: np.ndarray) -> List[Dict[str, Any]]:
|
||||
@@ -62,7 +95,7 @@ class Model:
|
||||
original_height, original_width = image.shape[:2]
|
||||
|
||||
try:
|
||||
# YOLOv8推理
|
||||
# YOLO 推理
|
||||
results = self.model(
|
||||
image,
|
||||
conf=self.conf_threshold,
|
||||
@@ -122,7 +155,7 @@ class Model:
|
||||
@property
|
||||
def applies_nms(self) -> bool:
|
||||
"""模型是否内部应用了 NMS"""
|
||||
# YOLOv8会自动应用 NMS
|
||||
# Ultralytics YOLO 会自动应用 NMS
|
||||
return True
|
||||
|
||||
def close(self):
|
||||
@@ -132,4 +165,5 @@ class Model:
|
||||
del self.model
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
print("YOLOv8模型已关闭")
|
||||
print(f"{self.model_name}模型已关闭")
|
||||
|
||||
@@ -1,56 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import sys
|
||||
import os
|
||||
|
||||
def attempt_load(weights, device=''):
|
||||
"""尝试加载YOLOv5模型"""
|
||||
# 加载模型
|
||||
model = torch.load(weights, map_location=device)
|
||||
|
||||
# 确定模型格式
|
||||
if isinstance(model, dict):
|
||||
if 'model' in model: # state_dict格式
|
||||
model = model['model']
|
||||
elif 'state_dict' in model: # state_dict格式
|
||||
model = model['state_dict']
|
||||
|
||||
# 如果是state_dict，则需要创建模型架构
|
||||
if isinstance(model, dict):
|
||||
print("警告：加载的是权重字典，尝试创建默认模型结构")
|
||||
from models.yolov5_model import YOLOv5
|
||||
model_arch = YOLOv5()
|
||||
model_arch.load_state_dict(model)
|
||||
model = model_arch
|
||||
|
||||
# 设置为评估模式
|
||||
if isinstance(model, nn.Module):
|
||||
model.eval()
|
||||
|
||||
# 检查是否有类别信息
|
||||
if not hasattr(model, 'names') or not model.names:
|
||||
print("模型没有类别信息，尝试加载默认类别")
|
||||
# 设置通用类别
|
||||
model.names = ['object']
|
||||
|
||||
return model
|
||||
|
||||
class YOLOv5:
|
||||
"""简化版YOLOv5模型结构，用于加载权重"""
|
||||
def __init__(self):
|
||||
super(YOLOv5, self).__init__()
|
||||
self.names = [] # 类别名称
|
||||
# 这里应该添加真实的网络结构
|
||||
# 但为了简单起见，我们只提供一个占位符
|
||||
# 在实际使用中，您应该实现完整的网络架构
|
||||
|
||||
def forward(self, x):
|
||||
# 这里应该是实际的前向传播逻辑
|
||||
# 这只是一个占位符
|
||||
raise NotImplementedError("这是一个占位符模型，请使用完整的YOLOv5模型实现")
|
||||
|
||||
def load_state_dict(self, state_dict):
|
||||
print("尝试加载模型权重")
|
||||
# 实际的权重加载逻辑
|
||||
# 这只是一个占位符
|
||||
return self
|
||||
@@ -65,9 +65,9 @@ public class VideoAnalysisService {
|
||||
@Autowired
|
||||
private com.ruoyi.video.mapper.InspectionTaskRecordMapper inspectionTaskRecordMapper;
|
||||
|
||||
// 检测器配置 - 使用容器名而不是localhost
|
||||
private static final String PYTHON_API_URL = "http://localhost:8000/api/detect/file";
|
||||
private static final String MODEL_NAME = "yolov8_detector";
|
||||
// 检测器配置 - 支持环境变量配置
|
||||
private static final String PYTHON_API_URL = System.getenv().getOrDefault("PYTHON_API_URL", "http://localhost:8000") + "/api/detect/file";
|
||||
private static final String MODEL_NAME = "smoke"; // 默认使用吸烟检测模型
|
||||
|
||||
/**
|
||||
* 分析视频并更新记录（同步调用）
|
||||
|
||||
@@ -1,4 +1,16 @@
|
||||
[
|
||||
{"name":"smoke","path":"libs/models/smoke","size":[640,640],"backend":"opencv"},
|
||||
{"name":"garbage","path":"libs/models/garbage","size":[640,640],"backend":"opencv"}
|
||||
{
|
||||
"name": "smoke",
|
||||
"pythonModelName": "smoke",
|
||||
"pythonApiUrl": "http://localhost:8000/api/detect/file",
|
||||
"size": [640, 640],
|
||||
"backend": "python"
|
||||
},
|
||||
{
|
||||
"name": "garbage",
|
||||
"pythonModelName": "garbage",
|
||||
"pythonApiUrl": "http://localhost:8000/api/detect/file",
|
||||
"size": [640, 640],
|
||||
"backend": "python"
|
||||
}
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user