Update app.py

app.py

@@ -1,14 +1,10 @@
-# app.py
 from fastapi import FastAPI, File, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
-import numpy as np
 import tensorflow as tf
-from tensorflow.lite.python.interpreter import Interpreter
-import os
+import numpy as np
+from tensorflow import keras
 import google.generativeai as genai
-import uvicorn
-from typing import Optional
-from pydantic import BaseModel
+import os
 
 app = FastAPI()
 
@@ -21,72 +17,68 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Load the TensorFlow Lite model
-interpreter = Interpreter(model_path="model.tflite")
-interpreter.allocate_tensors()
+# Configure Gemini API
+GEMINI_API_KEY = os.getenv('GEMINI_API_KEY', 'AIzaSyBx0A7BA-nKVZOiVn39JXzdGKgeGQqwAFg')
+genai.configure(api_key=GEMINI_API_KEY)
+gemini_model = genai.GenerativeModel('gemini-pro')
 
-# Get input and output details
-input_details = interpreter.get_input_details()
-output_details = interpreter.get_output_details()
+# Load the model
+model = keras.models.load_model('Image_classify.keras')
 
-# Define categories
+# Define categories and image dimensions
 data_cat = ['disposable cups', 'paper', 'plastic bottle']
 img_height, img_width = 224, 224
 
-# Configure Gemini API
-GEMINI_API_KEY = os.getenv('GEMINI_API_KEY', 'AIzaSyBx0A7BA-nKVZOiVn39JXzdGKgeGQqwAFg')
-genai.configure(api_key=GEMINI_API_KEY)
-
-# Initialize Gemini model
-gemini_model = genai.GenerativeModel('gemini-pro')
+def generate_recycling_insight(detected_object):
+    """Generate sustainability insights for detected objects"""
+    try:
+        prompt = f"""
+        You are a sustainability-focused AI. Analyze the {detected_object} (which is a solid dry waste) 
+        and generate the top three innovative, eco-friendly recommendations for repurposing it. Ensure each recommendation is:
+        - Give the Title of the recommendation
+        - Practical and easy to implement
+        - Environmentally beneficial
+        - Clearly explained in one or two concise sentences
+        """
+        
+        response = gemini_model.generate_content(prompt)
+        return response.text.strip()
+    except Exception as e:
+        return f"Error generating insight: {str(e)}"
 
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
     try:
+        # Read and preprocess the image
         contents = await file.read()
-        
-        # Preprocess the image
-        img = tf.image.decode_image(contents, channels=3)
-        img = tf.image.resize(img, [img_height, img_width])
-        img_bat = np.expand_dims(img, 0).astype(np.float32)
-
-        # Set input tensor
-        interpreter.set_tensor(input_details[0]['index'], img_bat)
-
-        # Run inference
-        interpreter.invoke()
+        image_load = tf.image.decode_image(contents, channels=3)
+        image_load = tf.image.resize(image_load, [img_height, img_width])
+        img_bat = tf.expand_dims(image_load, 0)
 
-        # Get the result
-        output_data = interpreter.get_tensor(output_details[0]['index'])
-        predicted_class = data_cat[np.argmax(output_data)]
-        confidence = float(np.max(output_data) * 100)
-
-        # Generate sustainability insights with Gemini API
-        prompt = f"""
-        You are a sustainability-focused AI. Analyze the {predicted_class} (solid dry waste) 
-        and generate the top three innovative, eco-friendly recommendations for repurposing it. 
-        Each recommendation should:
-        - Provide a title
-        - Be practical and easy to implement
-        - Be environmentally beneficial
-        - Include a one or two-sentence explanation
-        Format each recommendation with a clear title followed by the explanation on a new line.
-        """
+        # Perform prediction
+        predictions = model.predict(img_bat, verbose=0)
+        score = tf.nn.softmax(predictions[0])
+        confidence = float(np.max(score) * 100)
+        
+        if confidence < 45:
+            return {
+                "error": "Confidence too low to make a prediction",
+                "confidence": confidence
+            }
+        
+        # Get prediction and insights
+        predicted_class = data_cat[np.argmax(score)]
+        sustainability_insight = generate_recycling_insight(predicted_class)
         
-        try:
-            response = gemini_model.generate_content(prompt)
-            insights = response.text.strip()
-        except Exception as e:
-            insights = f"Error generating insights: {str(e)}"
-
         return {
             "class": predicted_class,
             "confidence": confidence,
-            "insights": insights
+            "insights": sustainability_insight
         }
     
     except Exception as e:
         return {"error": str(e)}
 
 if __name__ == "__main__":
+    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)