Update app.py

app.py

@@ -1,290 +1,307 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.responses import JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
-import tensorflow as tf
-import numpy as np
-from PIL import Image
-import io
-import logging
-import uvicorn
-import os
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Initialize FastAPI app
-app = FastAPI(
-    title="Waste Classification API",
-    description="API for classifying waste into categories: Glass, Metal, Organic, Paper, Plastic",
-    version="1.0.0",
-    docs_url="/",  # Swagger UI at root for easy access
-)
-
-# Add CORS middleware for web access
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Global variables - match your training exactly
-model = None
-# IMPORTANT: Your class order from training (alphabetical from image_dataset_from_directory)
-class_labels = ["glass", "metal", "organic", "paper", "plastic"]
-
-def load_model():
-    """Load the trained TensorFlow/Keras model"""
-    try:
-        # Try loading different formats in order of preference
-        model_files = [
-            'waste_model.keras',  # Keras format (recommended)
-            'waste_model.h5',     # H5 format
-            'best_model.keras'    # Checkpoint from training
-        ]
-        
-        model = None
-        for model_file in model_files:
-            if os.path.exists(model_file):
-                try:
-                    model = tf.keras.models.load_model(model_file)
-                    logger.info(f"Model loaded successfully from {model_file}")
-                    break
-                except Exception as e:
-                    logger.warning(f"Failed to load {model_file}: {e}")
-                    continue
-        
-        if model is None:
-            logger.error("No model file found. Creating dummy model for testing.")
-            # Create dummy model with same architecture for testing
-            model = tf.keras.Sequential([
-                tf.keras.layers.Rescaling(1./255, input_shape=(224, 224, 3)),
-                tf.keras.applications.MobileNetV2(input_shape=(224, 224, 3), include_top=False, weights='imagenet'),
-                tf.keras.layers.GlobalAveragePooling2D(),
-                tf.keras.layers.Dense(128, activation='relu'),
-                tf.keras.layers.Dropout(0.2),
-                tf.keras.layers.Dense(5, activation='softmax')
-            ])
-            logger.warning("Using dummy model - predictions will be random!")
-            
-        return model
-        
-    except Exception as e:
-        logger.error(f"Critical error loading model: {e}")
-        raise Exception(f"Model loading failed: {e}")
-
-def preprocess_image(image_data):
-    """
-    Preprocess image to match your training pipeline
-    """
-    try:
-        # Load image
-        image = Image.open(io.BytesIO(image_data)).convert('RGB')
-        
-        # Resize to match training (224, 224)
-        image = image.resize((224, 224), Image.BICUBIC)  # Match your training interpolation
-        
-        # Convert to numpy array
-        image_array = np.array(image, dtype=np.float32)
-        
-        # Add batch dimension
-        image_array = np.expand_dims(image_array, axis=0)
-        
-        # NOTE: Your model has Rescaling(1./255) as first layer, so no need to normalize here
-        # The model will handle normalization internally
-        
-        return image_array
-        
-    except Exception as e:
-        logger.error(f"Image preprocessing error: {e}")
-        raise HTTPException(status_code=400, detail=f"Image preprocessing failed: {e}")
-
-@app.on_event("startup")
-async def startup_event():
-    """Load model on startup"""
-    global model
-    try:
-        model = load_model()
-        logger.info("API startup complete")
-        
-        # Test model with dummy input
-        dummy_input = np.random.random((1, 224, 224, 3)).astype(np.float32)
-        _ = model.predict(dummy_input, verbose=0)
-        logger.info("Model test prediction successful")
-        
-    except Exception as e:
-        logger.error(f"Startup failed: {e}")
-        raise
-
-@app.get("/health")
-async def health_check():
-    """Health check endpoint"""
-    try:
-        # Quick model test
-        dummy_input = np.random.random((1, 224, 224, 3)).astype(np.float32)
-        prediction = model.predict(dummy_input, verbose=0)
-        model_working = prediction is not None
-        
-        return {
-            "status": "healthy",
-            "model_loaded": model is not None,
-            "model_working": model_working,
-            "classes": class_labels,
-            "input_shape": "(224, 224, 3)",
-            "model_type": "TensorFlow/Keras MobileNetV2"
-        }
-    except Exception as e:
-        return {
-            "status": "unhealthy",
-            "error": str(e),
-            "model_loaded": model is not None,
-            "classes": class_labels
-        }
-
-@app.post("/classify")
-async def classify_image(file: UploadFile = File(...)):
-    """
-    Main classification endpoint for ESP32
-    
-    Expected usage:
-    curl -X POST -F "file=@image.jpg" https://your-space-url.hf.space/classify
-    
-    Returns:
-        JSON: {"label": "plastic"} or {"error": "message"}
-    """
-    try:
-        # Validate file type
-        if not file.content_type or not file.content_type.startswith('image/'):
-            logger.warning(f"Invalid file type: {file.content_type}")
-            raise HTTPException(status_code=400, detail="File must be an image")
-        
-        # Read image data
-        image_data = await file.read()
-        if len(image_data) == 0:
-            raise HTTPException(status_code=400, detail="Empty image file")
-            
-        logger.info(f"Processing image: {file.filename}, size: {len(image_data)} bytes")
-        
-        # Preprocess image
-        processed_image = preprocess_image(image_data)
-        
-        # Make prediction
-        predictions = model.predict(processed_image, verbose=0)
-        predicted_class_index = np.argmax(predictions[0])
-        predicted_class = class_labels[predicted_class_index]
-        confidence = float(predictions[0][predicted_class_index])
-        
-        logger.info(f"Prediction: {predicted_class} (confidence: {confidence:.3f})")
-        
-        # Return simple response for ESP32 - match your ESP32 expectation exactly
-        return {"label": predicted_class.capitalize()}  # Capitalize to match your ESP32 labels
-        
-    except HTTPException:
-        raise
-    except Exception as e:
-        logger.error(f"Classification error: {str(e)}")
-        return JSONResponse(
-            status_code=500, 
-            content={"error": f"Classification failed: {str(e)}"}
-        )
-
-@app.post("/classify/detailed")
-async def classify_detailed(file: UploadFile = File(...)):
-    """
-    Detailed classification endpoint with confidence scores
-    """
-    try:
-        # Validate file type
-        if not file.content_type or not file.content_type.startswith('image/'):
-            raise HTTPException(status_code=400, detail="File must be an image")
-        
-        # Read and process image
-        image_data = await file.read()
-        processed_image = preprocess_image(image_data)
-        
-        # Make prediction with full details
-        predictions = model.predict(processed_image, verbose=0)
-        predicted_class_index = np.argmax(predictions[0])
-        predicted_class = class_labels[predicted_class_index]
-        confidence = float(predictions[0][predicted_class_index])
-        
-        # Get all class probabilities
-        all_probs = {
-            class_labels[i].capitalize(): round(float(predictions[0][i]) * 100, 2) 
-            for i in range(len(class_labels))
-        }
-        
-        return {
-            "label": predicted_class.capitalize(),
-            "confidence": round(confidence * 100, 2),
-            "all_probabilities": all_probs,
-            "model_info": {
-                "architecture": "MobileNetV2",
-                "input_size": "224x224",
-                "classes": len(class_labels)
-            },
-            "status": "success"
-        }
-    
-    except Exception as e:
-        logger.error(f"Detailed classification error: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Classification failed: {str(e)}")
-
-@app.get("/info")
-async def get_info():
-    """API information endpoint"""
-    return {
-        "api_name": "Waste Classification API",
-        "version": "1.0.0",
-        "model": {
-            "architecture": "MobileNetV2 + Custom Head",
-            "framework": "TensorFlow/Keras",
-            "input_size": "224x224x3",
-            "preprocessing": "RGB, Resize, Rescaling (internal)"
-        },
-        "classes": [label.capitalize() for label in class_labels],
-        "endpoints": {
-            "/classify": "POST - Main classification endpoint (returns simple label)",
-            "/classify/detailed": "POST - Detailed classification with confidence",
-            "/health": "GET - Health check",
-            "/info": "GET - API information"
-        },
-        "usage": {
-            "esp32": "POST image to /classify endpoint",
-            "curl_example": "curl -X POST -F 'file=@image.jpg' https://your-space-url.hf.space/classify"
-        }
-    }
-
-@app.post("/test")
-async def test_with_dummy():
-    """Test endpoint with dummy data for debugging"""
-    try:
-        # Create dummy image (random noise)
-        dummy_image = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
-        dummy_input = np.expand_dims(dummy_image.astype(np.float32), axis=0)
-        
-        # Make prediction
-        predictions = model.predict(dummy_input, verbose=0)
-        predicted_class_index = np.argmax(predictions[0])
-        predicted_class = class_labels[predicted_class_index]
-        
-        return {
-            "test_status": "success",
-            "predicted_class": predicted_class.capitalize(),
-            "confidence": float(predictions[0][predicted_class_index]),
-            "all_predictions": [float(p) for p in predictions[0]]
-        }
-    except Exception as e:
-        return {"test_status": "failed", "error": str(e)}
-
-if __name__ == "__main__":
-    # Run the FastAPI app
-    port = int(os.environ.get("PORT", 7860))
-    uvicorn.run(
-        app, 
-        host="0.0.0.0", 
-        port=port,
-        log_level="info"
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+import tensorflow as tf
+import numpy as np
+from PIL import Image
+import io
+import logging
+import uvicorn
+import os
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Initialize FastAPI app
+app = FastAPI(
+    title="Waste Classification API",
+    description="API for classifying waste into categories: Glass, Metal, Organic, Paper, Plastic",
+    version="1.0.0",
+    docs_url="/",  # Swagger UI at root for easy access
+)
+
+# Add CORS middleware for web access
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Global variables - match your training exactly
+model = None
+# IMPORTANT: Your class order from training (alphabetical from image_dataset_from_directory)
+class_labels = ["glass", "metal", "organic", "paper", "plastic"]
+
+def load_model():
+    """Load the trained TensorFlow/Keras model"""
+    try:
+        # Try loading different formats in order of preference
+        model_files = [
+            'waste_model.keras',  # Keras format (recommended)
+            'waste_model.h5',     # H5 format
+            'best_model.keras'    # Checkpoint from training
+        ]
+        
+        model = None
+        for model_file in model_files:
+            if os.path.exists(model_file):
+                try:
+                    model = tf.keras.models.load_model(model_file)
+                    logger.info(f"Model loaded successfully from {model_file}")
+                    break
+                except Exception as e:
+                    logger.warning(f"Failed to load {model_file}: {e}")
+                    continue
+        
+        if model is None:
+            logger.error("No model file found. Creating dummy model for testing.")
+            # Create dummy model with same architecture for testing
+            model = tf.keras.Sequential([
+                tf.keras.layers.Rescaling(1./255, input_shape=(224, 224, 3)),
+                tf.keras.applications.MobileNetV2(input_shape=(224, 224, 3), include_top=False, weights='imagenet'),
+                tf.keras.layers.GlobalAveragePooling2D(),
+                tf.keras.layers.Dense(128, activation='relu'),
+                tf.keras.layers.Dropout(0.2),
+                tf.keras.layers.Dense(5, activation='softmax')
+            ])
+            logger.warning("Using dummy model - predictions will be random!")
+            
+        return model
+        
+    except Exception as e:
+        logger.error(f"Critical error loading model: {e}")
+        raise Exception(f"Model loading failed: {e}")
+
+def preprocess_image(image_data):
+    """
+    Preprocess image to match training pipeline:
+    - Crop ROI from ESP32 frame (400x296)
+    - Resize to 224x224
+    - Convert to numpy array, add batch dim
+    """
+    try:
+        # Load image
+        image = Image.open(io.BytesIO(image_data)).convert('RGB')
+        w, h = image.size  # (width=400, height=296 expected)
+
+        # Define ROI (x1:x2, y1:y2), inclusive indices
+        X1, X2 = 90, 280
+        Y1, Y2 = 5, 205
+
+        # Clamp to actual image bounds
+        X1 = max(0, min(w - 1, X1))
+        X2 = max(0, min(w - 1, X2))
+        Y1 = max(0, min(h - 1, Y1))
+        Y2 = max(0, min(h - 1, Y2))
+
+        # Pillow crop box is (left, upper, right, lower) with right/lower EXCLUSIVE
+        crop_box = (X1, Y1, X2 + 1, Y2 + 1)
+        image = image.crop(crop_box)
+
+        # Resize to model input size
+        image = image.resize((224, 224), Image.BICUBIC)
+
+        # Convert to numpy array
+        image_array = np.array(image, dtype=np.float32)
+
+        # Add batch dimension
+        image_array = np.expand_dims(image_array, axis=0)
+
+        # Model has Rescaling(1./255) layer, so no manual normalization
+        return image_array
+
+    except Exception as e:
+        logger.error(f"Image preprocessing error: {e}")
+        raise HTTPException(status_code=400, detail=f"Image preprocessing failed: {e}")
+
+
+@app.on_event("startup")
+async def startup_event():
+    """Load model on startup"""
+    global model
+    try:
+        model = load_model()
+        logger.info("API startup complete")
+        
+        # Test model with dummy input
+        dummy_input = np.random.random((1, 224, 224, 3)).astype(np.float32)
+        _ = model.predict(dummy_input, verbose=0)
+        logger.info("Model test prediction successful")
+        
+    except Exception as e:
+        logger.error(f"Startup failed: {e}")
+        raise
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    try:
+        # Quick model test
+        dummy_input = np.random.random((1, 224, 224, 3)).astype(np.float32)
+        prediction = model.predict(dummy_input, verbose=0)
+        model_working = prediction is not None
+        
+        return {
+            "status": "healthy",
+            "model_loaded": model is not None,
+            "model_working": model_working,
+            "classes": class_labels,
+            "input_shape": "(224, 224, 3)",
+            "model_type": "TensorFlow/Keras MobileNetV2"
+        }
+    except Exception as e:
+        return {
+            "status": "unhealthy",
+            "error": str(e),
+            "model_loaded": model is not None,
+            "classes": class_labels
+        }
+
+@app.post("/classify")
+async def classify_image(file: UploadFile = File(...)):
+    """
+    Main classification endpoint for ESP32
+    
+    Expected usage:
+    curl -X POST -F "file=@image.jpg" https://your-space-url.hf.space/classify
+    
+    Returns:
+        JSON: {"label": "plastic"} or {"error": "message"}
+    """
+    try:
+        # Validate file type
+        if not file.content_type or not file.content_type.startswith('image/'):
+            logger.warning(f"Invalid file type: {file.content_type}")
+            raise HTTPException(status_code=400, detail="File must be an image")
+        
+        # Read image data
+        image_data = await file.read()
+        if len(image_data) == 0:
+            raise HTTPException(status_code=400, detail="Empty image file")
+            
+        logger.info(f"Processing image: {file.filename}, size: {len(image_data)} bytes")
+        
+        # Preprocess image
+        processed_image = preprocess_image(image_data)
+        
+        # Make prediction
+        predictions = model.predict(processed_image, verbose=0)
+        predicted_class_index = np.argmax(predictions[0])
+        predicted_class = class_labels[predicted_class_index]
+        confidence = float(predictions[0][predicted_class_index])
+        
+        logger.info(f"Prediction: {predicted_class} (confidence: {confidence:.3f})")
+        
+        # Return simple response for ESP32 - match your ESP32 expectation exactly
+        return {"label": predicted_class.capitalize()}  # Capitalize to match your ESP32 labels
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Classification error: {str(e)}")
+        return JSONResponse(
+            status_code=500, 
+            content={"error": f"Classification failed: {str(e)}"}
+        )
+
+@app.post("/classify/detailed")
+async def classify_detailed(file: UploadFile = File(...)):
+    """
+    Detailed classification endpoint with confidence scores
+    """
+    try:
+        # Validate file type
+        if not file.content_type or not file.content_type.startswith('image/'):
+            raise HTTPException(status_code=400, detail="File must be an image")
+        
+        # Read and process image
+        image_data = await file.read()
+        processed_image = preprocess_image(image_data)
+        
+        # Make prediction with full details
+        predictions = model.predict(processed_image, verbose=0)
+        predicted_class_index = np.argmax(predictions[0])
+        predicted_class = class_labels[predicted_class_index]
+        confidence = float(predictions[0][predicted_class_index])
+        
+        # Get all class probabilities
+        all_probs = {
+            class_labels[i].capitalize(): round(float(predictions[0][i]) * 100, 2) 
+            for i in range(len(class_labels))
+        }
+        
+        return {
+            "label": predicted_class.capitalize(),
+            "confidence": round(confidence * 100, 2),
+            "all_probabilities": all_probs,
+            "model_info": {
+                "architecture": "MobileNetV2",
+                "input_size": "224x224",
+                "classes": len(class_labels)
+            },
+            "status": "success"
+        }
+    
+    except Exception as e:
+        logger.error(f"Detailed classification error: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Classification failed: {str(e)}")
+
+@app.get("/info")
+async def get_info():
+    """API information endpoint"""
+    return {
+        "api_name": "Waste Classification API",
+        "version": "1.0.0",
+        "model": {
+            "architecture": "MobileNetV2 + Custom Head",
+            "framework": "TensorFlow/Keras",
+            "input_size": "224x224x3",
+            "preprocessing": "RGB, Resize, Rescaling (internal)"
+        },
+        "classes": [label.capitalize() for label in class_labels],
+        "endpoints": {
+            "/classify": "POST - Main classification endpoint (returns simple label)",
+            "/classify/detailed": "POST - Detailed classification with confidence",
+            "/health": "GET - Health check",
+            "/info": "GET - API information"
+        },
+        "usage": {
+            "esp32": "POST image to /classify endpoint",
+            "curl_example": "curl -X POST -F 'file=@image.jpg' https://your-space-url.hf.space/classify"
+        }
+    }
+
+@app.post("/test")
+async def test_with_dummy():
+    """Test endpoint with dummy data for debugging"""
+    try:
+        # Create dummy image (random noise)
+        dummy_image = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
+        dummy_input = np.expand_dims(dummy_image.astype(np.float32), axis=0)
+        
+        # Make prediction
+        predictions = model.predict(dummy_input, verbose=0)
+        predicted_class_index = np.argmax(predictions[0])
+        predicted_class = class_labels[predicted_class_index]
+        
+        return {
+            "test_status": "success",
+            "predicted_class": predicted_class.capitalize(),
+            "confidence": float(predictions[0][predicted_class_index]),
+            "all_predictions": [float(p) for p in predictions[0]]
+        }
+    except Exception as e:
+        return {"test_status": "failed", "error": str(e)}
+
+if __name__ == "__main__":
+    # Run the FastAPI app
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run(
+        app, 
+        host="0.0.0.0", 
+        port=port,
+        log_level="info"
     )