Add logic for app.py

app.py

@@ -1,7 +1,40 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, File, UploadFile
+import tensorflow as tf
+import numpy as np
+from PIL import Image
+from io import BytesIO
+from fastapi.responses import JSONResponse
 
+# Load your trained model (make sure it's available in the working directory)
+model = tf.keras.models.load_model('model.h5')
+
+# Class names for predictions (modify if necessary)
+class_names = ['Glass', 'Metal', 'Paperboard', 'Plastic-Polystyrene', 'Plastic-Regular']
+
+# Create FastAPI app
 app = FastAPI()
 
-@app.get("/")
-def greet_json():
-    return {"Hello": "World!"}
+# Preprocessing the image (resize, normalize, etc.)
+def preprocess_image(image_file):
+    image = Image.open(image_file)
+    image = image.resize((240, 240))  # Resize image to match model input
+    img_array = np.array(image)  # Convert to numpy array
+    img_array = img_array.astype(np.float32) / 255.0  # Normalize
+    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
+    return img_array
+
+@app.post("/predict")
+async def predict(file: UploadFile = File(...)):
+    try:
+        img_array = preprocess_image(file.file)
+        predictions = model.predict(img_array)
+        predicted_class_idx = np.argmax(predictions, axis=1)[0]
+        predicted_class = class_names[predicted_class_idx]
+        return JSONResponse(content={"prediction": predicted_class})
+    except Exception as e:
+        return JSONResponse(content={"error": str(e)}, status_code=400)
+
+# If you want to manually run FastAPI (though Hugging Face will typically do this)
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)