Upload api_complete.py

api_complete.py

@@ -494,223 +494,75 @@ async def load_all_models():
             
             # Analyze the structure to understand the architecture
             if isinstance(crop_model_data, dict):
-                logger.info(f"Crop model keys: {list(crop_model_data.keys())[:10]}...")  # Show first 10 keys
+                # Use standard EfficientNet-B0 from torchvision to match training
+                from torchvision import models as tv_models
                 
-                # Show ALL classifier keys to understand the structure
-                classifier_keys = [k for k in crop_model_data.keys() if 'classifier' in k]
-                logger.info(f"📋 Classifier keys in checkpoint: {classifier_keys}")
-                logger.info(f"📋 Total checkpoint keys: {len(crop_model_data.keys())}")
-                
-                # Create the EXACT architecture matching checkpoint
-                class CustomEfficientNet(torch.nn.Module):
-                    def __init__(self, num_classes):
-                        super().__init__()
-                        # EXACT architecture matching your state dict keys
-                        self.features = torch.nn.Sequential(
-                            torch.nn.Conv2d(3, 32, 3, padding=1),    # features.0
-                            torch.nn.BatchNorm2d(32),                # features.1
-                            torch.nn.ReLU(inplace=True),             # features.2
-                            torch.nn.MaxPool2d(2),                   # features.3
-                            torch.nn.Conv2d(32, 64, 3, padding=1),   # features.4
-                            torch.nn.BatchNorm2d(64),                # features.5
-                            torch.nn.ReLU(inplace=True),             # features.6
-                            torch.nn.MaxPool2d(2),                   # features.7
-                            torch.nn.Conv2d(64, 128, 3, padding=1),  # features.8
-                            torch.nn.BatchNorm2d(128),               # features.9
-                            torch.nn.ReLU(inplace=True),             # features.10
-                            torch.nn.MaxPool2d(2),                   # features.11
-                            torch.nn.Conv2d(128, 256, 3, padding=1), # features.12
-                            torch.nn.BatchNorm2d(256),               # features.13
-                            torch.nn.ReLU(inplace=True),             # features.14
-                            torch.nn.MaxPool2d(2),                   # features.15
-                            torch.nn.Conv2d(256, 512, 3, padding=1), # features.16
-                            torch.nn.BatchNorm2d(512),               # features.17
-                            torch.nn.ReLU(inplace=True),             # features.18
-                        )
-                        
-                        # Classifier matching checkpoint EXACTLY
-                        # Checkpoint has Linear layers at indices 3 and 6
-                        # ['classifier.3.weight', 'classifier.3.bias', 'classifier.6.weight', 'classifier.6.bias']
-                        self.classifier = torch.nn.Sequential(
-                            torch.nn.Dropout(0.5),                   # classifier.0 (no weights)
-                            torch.nn.ReLU(inplace=True),             # classifier.1 (no weights)  
-                            torch.nn.Dropout(0.5),                   # classifier.2 (no weights)
-                            torch.nn.Linear(25088, 1024),            # classifier.3 - HAS WEIGHTS ✅
-                            torch.nn.ReLU(inplace=True),             # classifier.4 (no weights)
-                            torch.nn.Dropout(0.5),                   # classifier.5 (no weights)
-                            torch.nn.Linear(1024, num_classes)       # classifier.6 - HAS WEIGHTS ✅
-                        )
-                    
-                    def forward(self, x):
-                        x = self.features(x)
-                        x = torch.flatten(x, 1)
-                        x = self.classifier(x)
-                        return x
-                
-                # Create models with correct number of classes
-                num_crop_classes = len(crop_classes['classes'])
-                num_disease_classes = len(disease_classes['classes'])
-                
-                logger.info(f"🎯 Creating crop model with {num_crop_classes} classes")
-                logger.info(f"🎯 Creating disease model with {num_disease_classes} classes")
-                
-                crop_model = CustomEfficientNet(num_crop_classes)
-                disease_model = CustomEfficientNet(num_disease_classes)
-                
-                # Load state dicts with strict=False to handle any mismatches
-                logger.info("📊 Loading crop model state_dict...")
-                crop_result = crop_model.load_state_dict(crop_model_data, strict=False)
-                crop_missing = list(crop_result.missing_keys) if hasattr(crop_result, 'missing_keys') else []
-                crop_unexpected = list(crop_result.unexpected_keys) if hasattr(crop_result, 'unexpected_keys') else []
-                
-                if crop_unexpected:
-                    logger.info(f"📊 Crop model unexpected keys: {crop_unexpected[:5]}...")
-                
-                logger.info("📊 Loading disease model state_dict...")
-                disease_result = disease_model.load_state_dict(disease_model_data, strict=False)
-                disease_missing = list(disease_result.missing_keys) if hasattr(disease_result, 'missing_keys') else []
-                disease_unexpected = list(disease_result.unexpected_keys) if hasattr(disease_result, 'unexpected_keys') else []
-                
-                # Filter out non-critical missing keys
-                non_critical_patterns = ['num_batches_tracked', 'running_mean', 'running_var']
-                
-                def filter_critical_keys(missing_keys):
-                    if not missing_keys:
-                        return []
-                    return [k for k in missing_keys if not any(pattern in k for pattern in non_critical_patterns)]
-                
-                crop_critical = filter_critical_keys(crop_missing)
-                disease_critical = filter_critical_keys(disease_missing)
-                
-                # Log missing keys
-                if crop_missing:
-                    logger.info(f"📊 Crop model missing keys: {len(crop_missing)} (filtering non-critical...)")
-                if disease_missing:
-                    logger.info(f"📊 Disease model missing keys: {len(disease_missing)} (filtering non-critical...)")
-                
-                # AUTO-FIX: Rebuild classifier if critical keys are missing
-                def fix_missing_classifier(model, missing_keys, model_name, num_classes):
-                    """Automatically fix missing classifier keys by rebuilding the layer"""
-                    classifier_missing = [k for k in missing_keys if 'classifier' in k]
-                    
-                    if classifier_missing:
-                        logger.warning(f"⚠️ {model_name} missing classifier keys: {classifier_missing}")
-                        logger.info(f"🔧 Auto-fixing: Rebuilding classifier for {num_classes} classes...")
-                        
-                        # Rebuild classifier matching checkpoint architecture EXACTLY
-                        # Checkpoint has Linear layers at indices 3 and 6
-                        model.classifier = torch.nn.Sequential(
-                            torch.nn.Dropout(0.5),                   # classifier.0 (no weights)
-                            torch.nn.ReLU(inplace=True),             # classifier.1 (no weights)
-                            torch.nn.Dropout(0.5),                   # classifier.2 (no weights)
-                            torch.nn.Linear(25088, 1024),            # classifier.3 - HAS WEIGHTS ✅
-                            torch.nn.ReLU(inplace=True),             # classifier.4 (no weights)
-                            torch.nn.Dropout(0.5),                   # classifier.5 (no weights)
-                            torch.nn.Linear(1024, num_classes)       # classifier.6 - HAS WEIGHTS ✅
-                        )
+                def load_efficientnet(state_dict, num_classes, model_name):
+                    try:
+                        logger.info(f"🏗️ Building EfficientNet-B0 for {model_name}...")
+                        # 1. Init standard model
+                        model = tv_models.efficientnet_b0(weights=None)
                         
-                        logger.info(f"✅ {model_name} classifier rebuilt with output: {num_classes} classes")
+                        # 2. Modify classifier to match num_classes (1280 -> num_classes)
+                        # EfficientNet-B0 classifier is: Sequential(Dropout, Linear(1280, 1000))
+                        # We need to change the Linear layer at index 1
+                        in_features = model.classifier[1].in_features
+                        model.classifier[1] = torch.nn.Linear(in_features, num_classes)
                         
-                        # Try to reload state_dict after fixing
+                        # 3. Load weights
+                        msg = model.load_state_dict(state_dict, strict=True)
+                        logger.info(f"✅ {model_name} loaded successfully (Strict=True)")
+                        return model
+                    except Exception as e:
+                        logger.warning(f"⚠️ Strict loading failed for {model_name}: {e}")
+                        logger.info("🔄 Retrying with strict=False...")
                         try:
-                            model.load_state_dict(crop_model_data if 'crop' in model_name.lower() else disease_model_data, strict=False)
-                            logger.info(f"✅ {model_name} state_dict reloaded after classifier fix")
-                        except Exception as reload_error:
-                            logger.warning(f"⚠️ Could not reload state_dict: {reload_error}")
-                    
-                    return model
+                            model.load_state_dict(state_dict, strict=False)
+                            logger.info(f"✅ {model_name} loaded (Strict=False)")
+                            return model
+                        except Exception as e2:
+                            logger.error(f"❌ Failed to load {model_name}: {e2}")
+                            raise e2
+
+                # Create and load models
+                num_crop_classes = len(crop_classes['classes'])
+                num_disease_classes = len(disease_classes['classes'])
                 
-                # Fix missing classifier keys automatically
-                if crop_critical:
-                    crop_model = fix_missing_classifier(crop_model, crop_critical, "Crop model", num_crop_classes)
-                    # Re-check for remaining critical keys
-                    crop_critical_remaining = [k for k in crop_critical if 'classifier' not in k]
-                    if crop_critical_remaining:
-                        logger.warning(f"⚠️ Crop model still missing keys: {crop_critical_remaining}")
+                logger.info(f"🎯 Loading Crop Model ({num_crop_classes} classes)...")
+                crop_model = load_efficientnet(crop_model_data, num_crop_classes, "Crop Model")
                 
-                if disease_critical:
-                    disease_model = fix_missing_classifier(disease_model, disease_critical, "Disease model", num_disease_classes)
-                    # Re-check for remaining critical keys
-                    disease_critical_remaining = [k for k in disease_critical if 'classifier' not in k]
-                    if disease_critical_remaining:
-                        logger.warning(f"⚠️ Disease model still missing keys: {disease_critical_remaining}")
+                logger.info(f"🎯 Loading Disease Model ({num_disease_classes} classes)...")
+                disease_model = load_efficientnet(disease_model_data, num_disease_classes, "Disease Model")
                 
-                # Move models to device with optimal dtype
+                # Move to device
                 vision_dtype = torch.float16 if vision_device == "cuda" else torch.float32
-                logger.info(f"📊 Vision models dtype: {vision_dtype}")
-                
                 crop_model.to(vision_device, dtype=vision_dtype)
                 disease_model.to(vision_device, dtype=vision_dtype)
                 
-                # Set to evaluation mode
                 crop_model.eval()
                 disease_model.eval()
-                
-                logger.info(f"✅ Crop model moved to {vision_device} ({vision_dtype})")
-                logger.info(f"✅ Disease model moved to {vision_device} ({vision_dtype})")
-                
-                # Log final classifier shapes
-                crop_classifier_shape = list(crop_model.classifier[-1].weight.shape)
-                disease_classifier_shape = list(disease_model.classifier[-1].weight.shape)
-                logger.info(f"📊 Crop classifier output shape: {crop_classifier_shape}")
-                logger.info(f"📊 Disease classifier output shape: {disease_classifier_shape}")
-                
+            
             else:
-                # If they're complete model objects, use them directly
+                # If they are already model objects (legacy support)
                 crop_model = crop_model_data
                 disease_model = disease_model_data
-                if hasattr(crop_model, 'eval'):
-                    crop_model.eval()
-                if hasattr(disease_model, 'eval'):
-                    disease_model.eval()
+                crop_model.to(vision_device)
+                disease_model.to(vision_device)
+                crop_model.eval()
+                disease_model.eval()
+                    
+                logger.info(f"✅ Vision models loaded successfully on {vision_device}!")
                 
-            logger.info(f"✅ Vision models loaded successfully on {vision_device}!")
-            logger.info(f"🎯 Crop: {len(crop_classes.get('classes', []))} classes")
-            logger.info(f"🎯 Disease: {len(disease_classes.get('classes', []))} classes")
-            
         except Exception as e:
             logger.error(f"❌ Failed to load vision models: {e}")
-            
-            # Try CPU fallback if GPU failed
-            if vision_device == "cuda":
-                try:
-                    logger.warning("⚠️ GPU loading failed, trying CPU fallback...")
-                    vision_device = "cpu"
-                    
-                    # Retry loading on CPU
-                    crop_model_data = torch.load(crop_path / "best_model.pth", map_location='cpu')
-                    disease_model_data = torch.load(disease_path / "best_model.pth", map_location='cpu')
-                    
-                    crop_model = CustomEfficientNet(len(crop_classes['classes']))
-                    disease_model = CustomEfficientNet(len(disease_classes['classes']))
-                    
-                    crop_model.load_state_dict(crop_model_data, strict=False)
-                    disease_model.load_state_dict(disease_model_data, strict=False)
-                    
-                    crop_model.to(vision_device, dtype=torch.float32)
-                    disease_model.to(vision_device, dtype=torch.float32)
-                    crop_model.eval()
-                    disease_model.eval()
-                    
-                    logger.info(f"✅ Vision models loaded on CPU fallback")
-                except Exception as fallback_error:
-                    logger.error(f"❌ CPU fallback also failed: {fallback_error}")
-                    import traceback
-                    logger.error(f"Full error: {traceback.format_exc()}")
-                    # Create dummy models for testing
-                    crop_model = None
-                    disease_model = None
-                    crop_classes = {'classes': ['tomato', 'potato', 'wheat', 'rice']}
-                    disease_classes = {'classes': ['healthy', 'early_blight', 'late_blight', 'leaf_spot']}
-            else:
-                import traceback
-                logger.error(f"Full error: {traceback.format_exc()}")
-                # Create dummy models for testing
-                crop_model = None
-                disease_model = None
-                crop_classes = {'classes': ['tomato', 'potato', 'wheat', 'rice']}
-                disease_classes = {'classes': ['healthy', 'early_blight', 'late_blight', 'leaf_spot']}
-            
+            import traceback
+            logger.error(f"Full error: {traceback.format_exc()}")
+            # Create dummy models for testing
+            crop_model = None
+            disease_model = None
+            crop_classes = {'classes': ['tomato', 'potato', 'wheat', 'rice']}
+            disease_classes = {'classes': ['healthy', 'early_blight', 'late_blight', 'leaf_spot']}
+
         # Define image transforms - match training resolution
         # 100352 / 512 = 196, so sqrt(196) = 14, meaning 14x14 feature map
         # This suggests input should be smaller to get 25088 features
@@ -752,7 +604,8 @@ async def load_all_models():
                 
                 # Load model with joblib
                 market_model = joblib.load(model_path)
-                logger.info(f"Loaded {model_path.name} with joblib")
+                from pathlib import Path
+                logger.info(f"Loaded {Path(model_path).name} with joblib")
                 
                 # Load encoders with joblib
                 encoders = joblib.load(encoders_path)
@@ -1355,7 +1208,7 @@ async def agricultural_chat(
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.post("/image-diagnosis")
-async def image_diagnosis(image_file: UploadFile = File(...)):
+async def image_diagnosis(image_file: UploadFile = File(...), language: str = Form("en")):
     """Diagnose crop diseases from images"""
     try:
         if 'vision' not in models or models['vision'] is None:
@@ -1412,12 +1265,27 @@ async def image_diagnosis(image_file: UploadFile = File(...)):
                 
                 with torch.no_grad():
                     outputs = disease_model(image_tensor)
-                    probabilities = torch.nn.functional.softmax(outputs, dim=1)
-                    predicted_class = torch.argmax(probabilities, dim=1).item()
-                    confidence = probabilities[0][predicted_class].item()
+                    # Get Top 3
+                    probs = torch.nn.functional.softmax(outputs, dim=1)
+                    top3_prob, top3_idx = torch.topk(probs, min(3, len(disease_classes['classes']))) # Ensure we don't ask for more than available classes
+                    
+                    # Primary prediction (Top 1)
+                    top1_idx = top3_idx[0][0].item()
+                    confidence = top3_prob[0][0].item()
+                    disease_name = disease_classes['classes'][top1_idx]
+                    
+                    # Top 3 List
+                    predictions = []
+                    for i in range(top3_idx.shape[1]):
+                        idx = top3_idx[0][i].item()
+                        prob = top3_prob[0][i].item()
+                        predictions.append({
+                            "disease": disease_classes['classes'][idx],
+                            "confidence": round(prob * 100, 2)
+                        })
                 
-                disease_name = disease_classes['classes'][predicted_class]
                 logger.info(f"Disease prediction: {disease_name} (confidence: {confidence:.3f})")
+                logger.info(f"Top 3 predictions: {predictions}")
             except Exception as model_error:
                 logger.error(f"Model inference error: {model_error}")
                 # Fallback when model inference fails
@@ -1431,32 +1299,58 @@ async def image_diagnosis(image_file: UploadFile = File(...)):
             confidence = random.uniform(0.7, 0.95)
         
         # Generate treatment recommendations
-        treatments = {
-            "Early_Blight": [
-                "Remove infected leaves immediately",
-                "Apply Mancozeb fungicide (2.5g/L water)",
-                "Spray Trichoderma solution weekly",
-                "Improve air circulation"
-            ],
-            "Late_Blight": [
-                "Apply Copper oxychloride spray",
-                "Remove infected plants",
-                "Avoid overhead watering",
-                "Use resistant varieties"
-            ]
-        }
+        try:
+            import json
+            info_path = Path(__file__).parent / "disease_info.json"
+            if info_path.exists():
+                with open(info_path, 'r') as f:
+                    disease_info = json.load(f)
+            else:
+                disease_info = {}
+        except Exception:
+            disease_info = {}
+
+        info = disease_info.get(disease_name, {})
         
-        treatment = treatments.get(disease_name, [
+        # Default fallback if disease not in JSON
+        default_treatment = [
             "Consult agricultural expert",
             "Apply appropriate fungicide",
             "Maintain proper plant hygiene"
-        ])
+        ]
         
+        treatment = info.get("treatment", default_treatment)
+        cause = info.get("cause", "Fungal or bacterial infection caused by environmental conditions")
+        prevention = info.get("prevention", "Use resistant varieties and practice crop rotation")
+
+        # --- TRANSLATION LOGIC ---
+        if language != "en":
+            try:
+                from deep_translator import GoogleTranslator
+                translator = GoogleTranslator(source='auto', target=language)
+                
+                # Translate Cause
+                cause = translator.translate(cause)
+                
+                # Translate Prevention
+                prevention = translator.translate(prevention)
+                
+                # Translate Treatment List
+                translated_treatments = []
+                for t in treatment:
+                    translated_treatments.append(translator.translate(t))
+                treatment = translated_treatments
+                
+            except Exception as trans_e:
+                logger.error(f"Translation failed: {trans_e}")
+
         return {
             "disease": disease_name,
             "confidence": round(confidence * 100, 2),
             "treatment": treatment,
-            "cause": "Fungal infection caused by environmental conditions"
+            "cause": cause,
+            "prevention": prevention,
+            "top_3_predictions": predictions
         }
         
     except Exception as e:
@@ -1691,5 +1585,5 @@ async def text_to_speech(
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=8000)