HVC-Tunnel-Multimodal/Infrared-ai-research-lab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

feat: add device parameter support for GPU/CPU switching in inference API

Browse Source 
- InfraredDetector.detect() 和 detect_and_draw() 新增 device 参数
- TempFeatureExtractor.detect_and_extract() 新增 device 参数
- FaultDiagnoser.diagnose() 和 batch_diagnose() 新增 device 参数
- batch_extract_features() 新增 device 参数
- 命令行 --device 选项支持指定推理设备
- detect_and_draw() 返回值改为 (image, n_detections) 避免重复检测
This commit is contained in:
SkyLi
2026-03-27 09:10:57 +08:00
parent 9a5f6fe473
commit febd9eea8d
3 changed files with 32 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/detect.py
+12 -7
View File
@@ -34,13 +34,15 @@ class InfraredDetector:
print(f"Model loaded from: {model_path}")
def detect(self, image: Union[str, np.ndarray],
conf: Optional[float] = None) -> List[Dict]:
conf: Optional[float] = None,
device: Optional[str] = None) -> List[Dict]:
"""
对图像进行目标检测
Args:
image: 图像路径或图像数组
conf: 置信度阈值（覆盖默认值）
device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU
Returns:
检测结果列表
@@ -49,7 +51,7 @@ class InfraredDetector:
raise ValueError("Model not loaded. Please load a model first.")
threshold = conf if conf is not None else self.conf
results = self.model(image, conf=threshold)
results = self.model(image, conf=threshold, device=device)
detections = []
for result in results:
@@ -75,7 +77,8 @@ class InfraredDetector:
def detect_and_draw(self, image: Union[str, np.ndarray],
output_path: Optional[str] = None,
conf: Optional[float] = None) -> np.ndarray:
conf: Optional[float] = None,
device: Optional[str] = None) -> np.ndarray:
"""
检测并绘制结果
@@ -83,6 +86,7 @@ class InfraredDetector:
image: 图像路径或图像数组
output_path: 输出图像路径
conf: 置信度阈值
device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU
Returns:
绘制了检测结果的图像
@@ -92,7 +96,7 @@ class InfraredDetector:
else:
img = image.copy()
detections = self.detect(img, conf)
detections = self.detect(img, conf, device=device)
for det in detections:
x1, y1, x2, y2 = map(int, det['bbox'])
@@ -110,7 +114,7 @@ class InfraredDetector:
cv2.imwrite(output_path, img)
print(f"Result saved to: {output_path}")
return img
return img, len(detections)
def train_yolo_model(data_yaml: str,
@@ -185,6 +189,7 @@ if __name__ == '__main__':
detect_parser.add_argument('--image', type=str, required=True, help='Image path')
detect_parser.add_argument('--output', type=str, help='Output image path')
detect_parser.add_argument('--conf', type=float, default=0.25, help='Confidence threshold')
detect_parser.add_argument('--device', type=str, default=None, help='Device (e.g., "0" for GPU, "cpu" for CPU)')
# Train command
train_parser = subparsers.add_parser('train', help='Train YOLO model')
@@ -198,8 +203,8 @@ if __name__ == '__main__':
if args.command == 'detect':
detector = InfraredDetector(args.model, args.conf)
img = detector.detect_and_draw(args.image, args.output)
print(f"Detected {len(detector.detect(args.image))} objects")
img, n_detections = detector.detect_and_draw(args.image, args.output, device=args.device)
print(f"Detected {n_detections} objects")
elif args.command == 'train':
best_path = train_yolo_model(
src/diagnose.py
+10 -5
View File
@@ -29,12 +29,14 @@ class FaultDiagnoser:
self.extractor = TempFeatureExtractor(model_path)
self.temp_threshold = temp_threshold or {}
def diagnose(self, image: Union[str, np.ndarray]) -> Dict:
def diagnose(self, image: Union[str, np.ndarray],
device: Optional[str] = None) -> Dict:
"""
对图像进行故障诊断
Args:
image: 图像路径或数组
device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU
Returns:
诊断结果
@@ -47,7 +49,7 @@ class FaultDiagnoser:
}
if self.detector and self.detector.model:
detections = self.detector.detect(image)
detections = self.detector.detect(image, device=device)
result['detections'] = detections
import cv2
@@ -148,13 +150,15 @@ class FaultDiagnoser:
}
def batch_diagnose(self, image_dir: Union[str, Path],
output_path: Union[str, Path]) -> None:
output_path: Union[str, Path],
device: Optional[str] = None) -> None:
"""
批量诊断
Args:
image_dir: 图像目录
output_path: 输出JSON路径
device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU
"""
image_dir = Path(image_dir)
results = []
@@ -164,7 +168,7 @@ class FaultDiagnoser:
continue
try:
result = self.diagnose(str(img_path))
result = self.diagnose(str(img_path), device=device)
result['image'] = img_path.name
results.append(result)
except Exception as e:
@@ -192,6 +196,7 @@ if __name__ == '__main__':
parser.add_argument('--image-dir', type=str, required=True, help='Image directory')
parser.add_argument('--output', type=str, required=True, help='Output JSON path')
parser.add_argument('--threshold', type=str, help='Threshold config JSON')
parser.add_argument('--device', type=str, default=None, help='Device (e.g., "0" for GPU, "cpu" for CPU)')
args = parser.parse_args()
@@ -200,4 +205,4 @@ if __name__ == '__main__':
threshold = load_threshold_config(args.threshold)
diagnoser = FaultDiagnoser(args.model, threshold)
diagnoser.batch_diagnose(args.image_dir, args.output)
diagnoser.batch_diagnose(args.image_dir, args.output, device=args.device)
src/extract_temp.py
+10 -5
View File
@@ -33,13 +33,15 @@ class TempFeatureExtractor:
self.model = YOLO(model_path)
def detect_and_extract(self, image: np.ndarray,
conf: float = 0.25) -> List[Dict]:
conf: float = 0.25,
device: Optional[str] = None) -> List[Dict]:
"""
对图像进行检测并提取ROI温度特征
Args:
image: 输入红外图像
conf: 检测置信度阈值
device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU
Returns:
检测结果列表，每个元素包含检测框和温度特征
@@ -47,7 +49,7 @@ class TempFeatureExtractor:
if self.model is None:
raise ValueError("Model not loaded. Please provide a valid model path.")
results = self.model(image, conf=conf)
results = self.model(image, conf=conf, device=device)
detections = []
for result in results:
@@ -135,7 +137,8 @@ def load_infrared_image(image_path: str) -> np.ndarray:
def batch_extract_features(image_dir: Union[str, Path],
output_path: Union[str, Path],
model_path: Optional[str] = None,
conf: float = 0.25) -> None:
conf: float = 0.25,
device: Optional[str] = None) -> None:
"""
批量处理图像并提取温度特征
@@ -144,6 +147,7 @@ def batch_extract_features(image_dir: Union[str, Path],
output_path: 输出JSON文件路径
model_path: YOLO模型路径
conf: 检测置信度阈值
device: 运行设备，例如 '0' 表示GPU 0，'cpu' 表示CPU
"""
image_dir = Path(image_dir)
extractor = TempFeatureExtractor(model_path)
@@ -155,7 +159,7 @@ def batch_extract_features(image_dir: Union[str, Path],
try:
image = load_infrared_image(str(img_path))
detections = extractor.detect_and_extract(image, conf=conf)
detections = extractor.detect_and_extract(image, conf=conf, device=device)
all_results.append({
'image': img_path.name,
@@ -179,6 +183,7 @@ if __name__ == '__main__':
parser.add_argument('--output', type=str, required=True, help='Output JSON path')
parser.add_argument('--model', type=str, default=None, help='YOLO model path')
parser.add_argument('--conf', type=float, default=0.25, help='Detection confidence threshold')
parser.add_argument('--device', type=str, default=None, help='Device (e.g., "0" for GPU, "cpu" for CPU)')
args = parser.parse_args()
batch_extract_features(args.image_dir, args.output, args.model, args.conf)
batch_extract_features(args.image_dir, args.output, args.model, args.conf, args.device)