Cloud110702 commited on
Commit
af3bb2f
·
verified ·
1 Parent(s): 43be5e0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -20
app.py CHANGED
@@ -2,7 +2,6 @@ from fastapi import FastAPI, File, UploadFile
2
  from fastapi.middleware.cors import CORSMiddleware
3
  import tensorflow as tf
4
  import numpy as np
5
- from tensorflow.lite.python.interpreter import Interpreter
6
  import google.generativeai as genai
7
  import os
8
 
@@ -22,13 +21,12 @@ GEMINI_API_KEY = os.getenv('GEMINI_API_KEY', 'AIzaSyBx0A7BA-nKVZOiVn39JXzdGKgeGQ
22
  genai.configure(api_key=GEMINI_API_KEY)
23
  gemini_model = genai.GenerativeModel('gemini-pro')
24
 
25
- # Load TFLite model
26
- interpreter = Interpreter(model_path="model.tflite")
27
- interpreter.allocate_tensors()
28
-
29
- # Get input and output details
30
- input_details = interpreter.get_input_details()
31
- output_details = interpreter.get_output_details()
32
 
33
  # Define categories and image dimensions
34
  data_cat = ['disposable cups', 'paper', 'plastic bottle']
@@ -61,17 +59,10 @@ async def predict(file: UploadFile = File(...)):
61
  image = tf.cast(image, tf.float32)
62
  image = tf.expand_dims(image, 0)
63
 
64
- # Set the input tensor
65
- interpreter.set_tensor(input_details[0]['index'], image)
66
-
67
- # Run inference
68
- interpreter.invoke()
69
-
70
- # Get the output tensor
71
- output_data = interpreter.get_tensor(output_details[0]['index'])
72
-
73
- # Calculate confidence and get prediction
74
- confidence = float(np.max(output_data) * 100)
75
 
76
  if confidence < 45:
77
  return {
@@ -79,7 +70,7 @@ async def predict(file: UploadFile = File(...)):
79
  "confidence": confidence
80
  }
81
 
82
- predicted_class = data_cat[np.argmax(output_data)]
83
  sustainability_insight = generate_recycling_insight(predicted_class)
84
 
85
  return {
 
2
  from fastapi.middleware.cors import CORSMiddleware
3
  import tensorflow as tf
4
  import numpy as np
 
5
  import google.generativeai as genai
6
  import os
7
 
 
21
  genai.configure(api_key=GEMINI_API_KEY)
22
  gemini_model = genai.GenerativeModel('gemini-pro')
23
 
24
+ # Load model with specific version handling
25
+ model = tf.keras.models.load_model(
26
+ 'Image_classify.keras',
27
+ custom_objects=None,
28
+ compile=False # Don't compile the model on load
29
+ )
 
30
 
31
  # Define categories and image dimensions
32
  data_cat = ['disposable cups', 'paper', 'plastic bottle']
 
59
  image = tf.cast(image, tf.float32)
60
  image = tf.expand_dims(image, 0)
61
 
62
+ # Make prediction
63
+ predictions = model.predict(image, verbose=0)
64
+ score = tf.nn.softmax(predictions[0])
65
+ confidence = float(np.max(score) * 100)
 
 
 
 
 
 
 
66
 
67
  if confidence < 45:
68
  return {
 
70
  "confidence": confidence
71
  }
72
 
73
+ predicted_class = data_cat[np.argmax(score)]
74
  sustainability_insight = generate_recycling_insight(predicted_class)
75
 
76
  return {