Update app.py

app.py

@@ -2,7 +2,7 @@ from fastapi import FastAPI, File, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 import tensorflow as tf
 import numpy as np
-from tensorflow import keras
+from tensorflow.lite.python.interpreter import Interpreter
 import google.generativeai as genai
 import os
 
@@ -22,8 +22,13 @@ GEMINI_API_KEY = os.getenv('GEMINI_API_KEY', 'AIzaSyBx0A7BA-nKVZOiVn39JXzdGKgeGQ
 genai.configure(api_key=GEMINI_API_KEY)
 gemini_model = genai.GenerativeModel('gemini-pro')
 
-# Load the model
-model = keras.models.load_model('Image_classify.keras')
+# Load TFLite model
+interpreter = Interpreter(model_path="model.tflite")
+interpreter.allocate_tensors()
+
+# Get input and output details
+input_details = interpreter.get_input_details()
+output_details = interpreter.get_output_details()
 
 # Define categories and image dimensions
 data_cat = ['disposable cups', 'paper', 'plastic bottle']
@@ -51,14 +56,22 @@ async def predict(file: UploadFile = File(...)):
     try:
         # Read and preprocess the image
         contents = await file.read()
-        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)
+        image = tf.image.decode_image(contents, channels=3)
+        image = tf.image.resize(image, [img_height, img_width])
+        image = tf.cast(image, tf.float32)
+        image = tf.expand_dims(image, 0)
+
+        # Set the input tensor
+        interpreter.set_tensor(input_details[0]['index'], image)
 
-        # Perform prediction
-        predictions = model.predict(img_bat, verbose=0)
-        score = tf.nn.softmax(predictions[0])
-        confidence = float(np.max(score) * 100)
+        # Run inference
+        interpreter.invoke()
+
+        # Get the output tensor
+        output_data = interpreter.get_tensor(output_details[0]['index'])
+        
+        # Calculate confidence and get prediction
+        confidence = float(np.max(output_data) * 100)
         
         if confidence < 45:
             return {
@@ -66,8 +79,7 @@ async def predict(file: UploadFile = File(...)):
                 "confidence": confidence
             }
         
-        # Get prediction and insights
-        predicted_class = data_cat[np.argmax(score)]
+        predicted_class = data_cat[np.argmax(output_data)]
         sustainability_insight = generate_recycling_insight(predicted_class)
         
         return {