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

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +24 -12
app.py CHANGED
@@ -2,7 +2,7 @@ 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 import keras
6
  import google.generativeai as genai
7
  import os
8
 
@@ -22,8 +22,13 @@ 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 the model
26
- model = keras.models.load_model('Image_classify.keras')
 
 
 
 
 
27
 
28
  # Define categories and image dimensions
29
  data_cat = ['disposable cups', 'paper', 'plastic bottle']
@@ -51,14 +56,22 @@ async def predict(file: UploadFile = File(...)):
51
  try:
52
  # Read and preprocess the image
53
  contents = await file.read()
54
- image_load = tf.image.decode_image(contents, channels=3)
55
- image_load = tf.image.resize(image_load, [img_height, img_width])
56
- img_bat = tf.expand_dims(image_load, 0)
 
 
 
 
57
 
58
- # Perform prediction
59
- predictions = model.predict(img_bat, verbose=0)
60
- score = tf.nn.softmax(predictions[0])
61
- confidence = float(np.max(score) * 100)
 
 
 
 
62
 
63
  if confidence < 45:
64
  return {
@@ -66,8 +79,7 @@ async def predict(file: UploadFile = File(...)):
66
  "confidence": confidence
67
  }
68
 
69
- # Get prediction and insights
70
- predicted_class = data_cat[np.argmax(score)]
71
  sustainability_insight = generate_recycling_insight(predicted_class)
72
 
73
  return {
 
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
  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']
 
56
  try:
57
  # Read and preprocess the image
58
  contents = await file.read()
59
+ image = tf.image.decode_image(contents, channels=3)
60
+ image = tf.image.resize(image, [img_height, img_width])
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
  "confidence": confidence
80
  }
81
 
82
+ predicted_class = data_cat[np.argmax(output_data)]
 
83
  sustainability_insight = generate_recycling_insight(predicted_class)
84
 
85
  return {