Create main.py

main.py

@@ -0,0 +1,110 @@
+import io
+import logging
+from contextlib import asynccontextmanager
+
+import numpy as np
+import tensorflow as tf
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from PIL import Image
+from skimage import transform
+
+# --- Configuration ---
+IMG_SIZE = 224
+IMG_MODEL_FILENAME = "vgg_model50.h5" # Make sure this matches your uploaded file
+CLASS_NAMES_IMG = ["Non-Autistic", "Autistic"] # Adjust if your VGG model output differs
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+ml_models = {} # Dictionary to hold loaded models
+
+# --- Model Loading Logic ---
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Load the ML model during startup
+    logger.info(f"Attempting to load image model: {IMG_MODEL_FILENAME}")
+    try:
+        ml_models['image_classifier'] = tf.keras.models.load_model(IMG_MODEL_FILENAME)
+        logger.info("Image model loaded successfully.")
+    except Exception as e:
+        logger.error(f"Error loading image model '{IMG_MODEL_FILENAME}': {e}")
+        ml_models['image_classifier'] = None # Indicate loading failure
+    yield
+    # Clean up the ML models and release the resources
+    ml_models.clear()
+    logger.info("Cleaned up models.")
+
+# --- FastAPI App ---
+app = FastAPI(lifespan=lifespan)
+
+# --- CORS Middleware ---
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"], # Allows all origins
+    allow_credentials=True,
+    allow_methods=["*"], # Allows all methods
+    allow_headers=["*"], # Allows all headers
+)
+
+# --- Image Preprocessing ---
+def preprocess_image(image_bytes: bytes):
+    """Loads image bytes, preprocesses, and prepares for VGG16."""
+    try:
+        img = Image.open(io.BytesIO(image_bytes)).convert('RGB') # Ensure 3 channels
+        np_image = np.array(img).astype('float32') / 255.0 # Normalize
+        np_image = transform.resize(np_image, (IMG_SIZE, IMG_SIZE, 3))
+        np_image = np.expand_dims(np_image, axis=0) # Add batch dimension
+        return np_image
+    except Exception as e:
+        logger.error(f"Error preprocessing image: {e}")
+        raise HTTPException(status_code=400, detail=f"Error processing image file: {e}")
+
+# --- Prediction Endpoint (Image) ---
+@app.post("/predict/")
+async def predict_image(image: UploadFile = File(...)):
+    """Receives an image file, preprocesses it, and returns the VGG16 prediction."""
+    if ml_models.get('image_classifier') is None:
+         logger.error("Image model is not loaded.")
+         raise HTTPException(status_code=500, detail="Image model could not be loaded")
+
+    logger.info(f"Received image file: {image.filename}")
+    image_bytes = await image.read()
+    if not image_bytes:
+        raise HTTPException(status_code=400, detail="No image data received")
+
+    # Preprocess the image
+    processed_image = preprocess_image(image_bytes)
+
+    # Make prediction
+    try:
+        predictions = ml_models['image_classifier'].predict(processed_image)
+        predicted_class_index = np.argmax(predictions[0]) # Get index of highest probability
+
+        # --- IMPORTANT ADJUSTMENT ---
+        # Your VGG notebook used sparse_categorical_crossentropy and flow_from_directory
+        # with classes=['non_autistic','autistic']. This means index 0 is 'non_autistic' and 1 is 'autistic'.
+        # However, the final Dense layer had 95 units (output = Dense(95, activation='softmax')(class1)).
+        # This seems like a mismatch. Assuming the binary classification was intended:
+        if predicted_class_index < len(CLASS_NAMES_IMG):
+             predicted_class_name = CLASS_NAMES_IMG[predicted_class_index]
+             # For binary classification, maybe return probability too?
+             # probability = float(predictions[0][predicted_class_index])
+        else:
+             # Handle unexpected index if the model output isn't binary as expected
+             predicted_class_name = "Unknown Prediction"
+             logger.warning(f"Predicted index {predicted_class_index} is out of bounds for CLASS_NAMES_IMG.")
+
+        logger.info(f"Prediction successful: {predicted_class_name}")
+        return {"prediction": predicted_class_name}
+        # If you want probability: return {"prediction": predicted_class_name, "probability": probability}
+
+    except Exception as e:
+        logger.error(f"Error during prediction: {e}")
+        raise HTTPException(status_code=500, detail=f"Prediction error: {e}")
+
+# --- Root Endpoint (Optional - for health check/info) ---
+@app.get("/")
+async def root():
+    return {"message": "Autism Image Classification API"}