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

- 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) 避免重复检测

src/detect.py

@@ -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)
+        img, n_detections = detector.detect_and_draw(args.image, args.output, device=args.device)
-        print(f"Detected {len(detector.detect(args.image))} objects")
+        print(f"Detected {n_detections} objects")
 
     elif args.command == 'train':
         best_path = train_yolo_model(

src/diagnose.py

@@ -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

@@ -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)