Update pest/source/src/pest.py

pest/source/src/pest.py

@@ -8,119 +8,61 @@ def detect_pest(model, image_data):
     try:
         import torch
         import torch.nn.functional as F
-        import random
         
         # 检查输入类型
-        if isinstance(model, dict) and 'model' in model and isinstance(image_data, dict) and 'image_path' in image_data:
-            model_obj = model['model']
-            image_path = image_data['image_path']
+        if not (isinstance(model, dict) and 'model' in model and isinstance(image_data, dict) and 'image_path' in image_data):
+            raise ValueError("Invalid input format for model or image_data")
+        
+        model_obj = model['model']
+        image_path = image_data['image_path']
+        
+        # 检查是否是真实模型
+        if model.get('is_mock', True) != False or model_obj != 'real_model_loaded':
+            raise ValueError("Model is not properly loaded or is in mock mode")
+        
+        # 真实模型检测
+        print("Using real model for detection")
+        # 从全局变量获取模型和transform
+        from .utils import _loaded_model, _model_transform, _loaded_image
+        
+        if _loaded_model is None or _model_transform is None:
+            raise RuntimeError("Model or transform not loaded")
+        
+        if _loaded_image is None:
+            raise RuntimeError("Image not loaded")
+        
+        # 预处理图像
+        input_tensor = _model_transform(_loaded_image).unsqueeze(0)
+        
+        # 模型推理
+        with torch.no_grad():
+            outputs = _loaded_model(input_tensor)
+            probabilities = F.softmax(outputs, dim=1)
+            confidence, predicted = torch.max(probabilities, 1)
             
-            # 检查是否是真实模型
-            if model.get('is_mock', True) == False and model_obj == 'real_model_loaded':
-                # 真实模型检测
-                print("Using real model for detection")
-                # 从全局变量获取模型和transform
-                from .utils import _loaded_model, _model_transform
-                
-                if _loaded_model is not None and _model_transform is not None:
-                    # 从全局变量获取图像
-                    from .utils import _loaded_image
-                    print(f"Loaded model: {_loaded_model is not None}")
-                    print(f"Loaded transform: {_model_transform is not None}")
-                    print(f"Loaded image: {_loaded_image is not None}")
-                    if _loaded_image is not None:
-                        # 预处理图像
-                        input_tensor = _model_transform(_loaded_image).unsqueeze(0)
-                        
-                        # 模型推理
-                        with torch.no_grad():
-                            outputs = _loaded_model(input_tensor)
-                            probabilities = F.softmax(outputs, dim=1)
-                            confidence, predicted = torch.max(probabilities, 1)
-                            
-                            # 获取前5个预测结果
-                            top5_prob, top5_indices = torch.topk(probabilities, 5)
-                            
-                            predictions = []
-                            for i in range(5):
-                                class_id = top5_indices[0][i].item()
-                                confidence_score = top5_prob[0][i].item()
-                                
-                                predictions.append({
-                                    'class_id': class_id,
-                                    'confidence': confidence_score
-                                })
-                            
-                            return {
-                                'predictions': predictions,
-                                'top_prediction': {
-                                    'class_id': predicted.item(),
-                                    'confidence': confidence.item()
-                                },
-                                'is_mock': False
-                            }
-                
-                # 如果真实模型不可用，回退到模拟
-                print("Real model not available, falling back to mock")
-                mock_predictions = []
-                for i in range(5):
-                    class_id = random.randint(1, 102)
-                    confidence = random.uniform(0.1, 0.9)
-                    mock_predictions.append({
-                        'class_id': class_id,
-                        'confidence': confidence
-                    })
-                
-                return {
-                    'predictions': mock_predictions,
-                    'top_prediction': {
-                        'class_id': mock_predictions[0]['class_id'],
-                        'confidence': mock_predictions[0]['confidence']
-                    },
-                    'is_mock': True
-                }
-            else:
-                # 模拟检测结果
-                print("Using mock model for detection")
-                mock_predictions = []
-                for i in range(5):
-                    class_id = random.randint(1, 102)  # IP102数据集有102个类别
-                    confidence = random.uniform(0.1, 0.9)
-                    mock_predictions.append({
-                        'class_id': class_id,
-                        'confidence': confidence
-                    })
-                
-                return {
-                    'predictions': mock_predictions,
-                    'top_prediction': {
-                        'class_id': mock_predictions[0]['class_id'],
-                        'confidence': mock_predictions[0]['confidence']
-                    },
-                    'is_mock': True
-                }
-        else:
-            # 简化实现，返回模拟结果
-            print("Using simplified mock detection")
-            mock_predictions = []
+            # 获取前5个预测结果
+            top5_prob, top5_indices = torch.topk(probabilities, 5)
+            
+            predictions = []
             for i in range(5):
-                class_id = random.randint(1, 102)
-                confidence = random.uniform(0.1, 0.9)
-                mock_predictions.append({
+                class_id = top5_indices[0][i].item()
+                confidence_score = top5_prob[0][i].item()
+                
+                predictions.append({
                     'class_id': class_id,
-                    'confidence': confidence
+                    'confidence': confidence_score
                 })
             
             return {
-                'predictions': mock_predictions,
+                'predictions': predictions,
                 'top_prediction': {
-                    'class_id': mock_predictions[0]['class_id'],
-                    'confidence': mock_predictions[0]['confidence']
+                    'class_id': predicted.item(),
+                    'confidence': confidence.item()
                 },
-                'is_mock': True
+                'is_mock': False
             }
     except Exception as e:
-        return f"Detection error: {str(e)}"
+        raise RuntimeError(f"Detection error: {str(e)}")
 
 def initialize_detector():
     """Initialize the pest detector (only when needed)."""