提高视频帧率
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user