Update main.py

main.py

@@ -5,19 +5,18 @@ from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import JSONResponse
 from io import BytesIO
 from PIL import Image
-from tensorflow.keras.preprocessing.image import img_to_array
-from tensorflow.keras.applications import resnet50
-from tensorflow.keras.applications.resnet50 import preprocess_input
-import uvicorn
+from tensorflow.keras.preprocessing.image import img_to_array, load_img
 import tensorflow_addons as tfa
-# Load the h5 model
-custom_objects = {
-    "Addons>CohenKappa": tfa.metrics.CohenKappa,
-}
+import uvicorn
 
 # Initialize FastAPI app
 app = FastAPI()
 
+# Register the custom object
+custom_objects = {
+    "Addons>CohenKappa": tfa.metrics.CohenKappa,
+}
+
 # Model and class information
 model_path = "model.h5"
 class_labels = {
@@ -27,12 +26,12 @@ class_labels = {
     3: "Apple___healthy",
     4: "Background_without_leaves",
     5: "Blueberry___healthy",
-    6: "Cherry___Powdery_mildew",
-    7: "Cherry___healthy",
-    8: "Corn___Cercospora_leaf_spot_Gray_leaf_spot",
-    9: "Corn___Common_rust",
-    10: "Corn___Northern_Leaf_Blight",
-    11: "Corn___healthy",
+    6: "Cherry_(including_sour)___Powdery_mildew",
+    7: "Cherry_(including_sour)___healthy",
+    8: "Corn_(maize)___Cercospora_leaf_spot Gray_leaf_spot",
+    9: "Corn_(maize)___Common_rust_",
+    10: "Corn_(maize)___Northern_Leaf_Blight",
+    11: "Corn_(maize)___healthy",
     12: "Grape___Black_rot",
     13: "Grape___Esca_(Black_Measles)",
     14: "Grape___Leaf_blight_(Isariopsis_Leaf_Spot)",
@@ -55,7 +54,7 @@ class_labels = {
     31: "Tomato___Late_blight",
     32: "Tomato___Leaf_Mold",
     33: "Tomato___Septoria_leaf_spot",
-    34: "Tomato___Spider_mites_Two-spotted_spider_mite",
+    34: "Tomato___Spider_mites Two-spotted_spider_mite",
     35: "Tomato___Target_Spot",
     36: "Tomato___Tomato_Yellow_Leaf_Curl_Virus",
     37: "Tomato___Tomato_mosaic_virus",
@@ -64,31 +63,38 @@ class_labels = {
 
 # Load the model if it exists
 if os.path.exists(model_path):
-    model = tf.keras.models.load_model('model.h5', custom_objects=custom_objects)
+    model = tf.keras.models.load_model(model_path, custom_objects=custom_objects)
     print("Model loaded successfully.")
 else:
     print(f"Model file not found at {model_path}. Please upload the model.")
 
-# Function to predict crop disease in an image and return the class name
-def predict_image(image_data):
-    try:
-        # Load the image from binary data
-        img = Image.open(BytesIO(image_data))
-        # Resize the image to the target size
-        img = img.resize((224, 224))
-        # Convert image to array format for the model
-        img_array = img_to_array(img)
-        img_array = np.expand_dims(img_array, axis=0)
-        img_array = preprocess_input(img_array)
+# Function to preprocess input image
+def preprocess_image(image_data, img_size=224):
+    img = Image.open(BytesIO(image_data))
+    img = img.resize((img_size, img_size))
+    img_array = img_to_array(img)
+    img_array = img_array / 255.0
+    img_array = np.expand_dims(img_array, axis=0)
+    return img_array
 
-        # Make prediction
-        prediction = model.predict(img_array)
-        predicted_class = np.argmax(prediction[0])
-        class_name = class_labels.get(predicted_class, "Unknown")  # Map to class name
-        return class_name
-    except Exception as e:
-        print("Prediction error:", e)
-        return "Error during prediction"
+# Predict function
+def predict_image(image_data):
+    preprocessed_image = preprocess_image(image_data)
+    predictions = model.predict(preprocessed_image)
+    class_idx = np.argmax(predictions, axis=1)[0]
+    confidence = predictions[0][class_idx]
+    class_label = class_labels.get(class_idx, "Unknown")
+    if class_label is None:
+        return {
+            "class_index": class_idx,
+            "class_label": "Skipped/Invalid",
+            "confidence": confidence,
+        }
+    return {
+        "class_index": class_idx,
+        "class_label": class_label,
+        "confidence": confidence,
+    }
 
 # Route for health check
 @app.get("/health")
@@ -99,12 +105,8 @@ async def api_health_check():
 @app.post("/predict")
 async def api_predict_image(file: UploadFile = File(...)):
     try:
-        # Read the image file as binary data
         image_data = await file.read()
-        
-        # Call the prediction function with the image data
         prediction = predict_image(image_data)
-        
         return JSONResponse(content={"prediction": prediction})
     except Exception as e:
         return JSONResponse(content={"error": str(e)})