From febd9eea8d2014896d7ef21be3cdfeb559857287 Mon Sep 17 00:00:00 2001 From: SkyLi Date: Fri, 27 Mar 2026 09:10:57 +0800 Subject: [PATCH] feat: add device parameter support for GPU/CPU switching in inference API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 | 19 ++++++++++++------- src/diagnose.py | 15 ++++++++++----- src/extract_temp.py | 15 ++++++++++----- 3 files changed, 32 insertions(+), 17 deletions(-) diff --git a/src/detect.py b/src/detect.py index 0458de2..a0f0ae8 100644 --- a/src/detect.py +++ b/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) - 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( diff --git a/src/diagnose.py b/src/diagnose.py index c2d2dcd..356df7f 100644 --- a/src/diagnose.py +++ b/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) diff --git a/src/extract_temp.py b/src/extract_temp.py index 833fabd..98ff5b0 100644 --- a/src/extract_temp.py +++ b/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)