Spaces:
Sleeping
Sleeping
Update app.py
Browse files
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
|
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
|
26 |
-
|
|
|
|
|
|
|
|
|
|
|
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 |
-
|
55 |
-
|
56 |
-
|
|
|
|
|
|
|
|
|
57 |
|
58 |
-
#
|
59 |
-
|
60 |
-
|
61 |
-
|
|
|
|
|
|
|
|
|
62 |
|
63 |
if confidence < 45:
|
64 |
return {
|
@@ -66,8 +79,7 @@ async def predict(file: UploadFile = File(...)):
|
|
66 |
"confidence": confidence
|
67 |
}
|
68 |
|
69 |
-
|
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 {
|