Update app.py

app.py

@@ -2,25 +2,33 @@ from fastapi import FastAPI, File, UploadFile, Request
 import tensorflow as tf
 import numpy as np
 from PIL import Image
-import cv2  # Make sure OpenCV is available for resizing
+import cv2
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
-#from slowapi import Limiter
-#from slowapi import Limiter
-#from slowapi.util import get_remote_address
-#from slowapi.errors import RateLimitExceeded
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
 
 # Load your trained model
 model = tf.keras.models.load_model('recyclebot.keras')
 
+# Load background removal model
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    "ZhengPeng7/BiRefNet", trust_remote_code=True
+)
+
+# Transform for the background removal model
+transform_image = transforms.Compose([
+    transforms.Resize((1024, 1024)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
 # Define class names for predictions (this should be the same as in your local code)
 CLASSES = ['Glass', 'Metal', 'Paperboard', 'Plastic-Polystyrene', 'Plastic-Regular']
 
 # Create FastAPI app
 app = FastAPI()
 
-#limiter = Limiter(key_func=get_remote_address)
-
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],  # Allow all origins (or specify specific origins)
@@ -29,10 +37,7 @@ app.add_middleware(
     allow_headers=["*"],  # Allow all headers
 )
 
-
-
-
-# Preprocessing the image (resize, reshape without normalization)
+# Preprocess the image (resize, reshape without normalization)
 def preprocess_image(image_file):
     # Load image using PIL
     image = Image.open(image_file)
@@ -48,22 +53,25 @@ def preprocess_image(image_file):
     
     return image
 
-    
+# Background removal function
+def remove_background(image):
+    image_size = image.size
+    input_images = transform_image(image).unsqueeze(0)
+    with torch.no_grad():
+        preds = birefnet(input_images)[-1].sigmoid()
+    pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+    mask = pred_pil.resize(image_size)
+    image.putalpha(mask)
+    return image
 
 @app.post("/predict")
-#@limiter.limit("10/minute")
-async def predict(file: UploadFile = File(...)):   #async def predict(request: Request, file: UploadFile = File(...)):
+async def predict(file: UploadFile = File(...)):
     try:
         img_array = preprocess_image(file.file)  # Preprocess the image
         prediction1 = model.predict(img_array)  # Get predictions
-
-        weight_1 = 0.6
-        weight_2 = 0.4
         
-        # Get the index of the highest probability class (like np.argmax on local machine)
         predicted_class_idx = np.argmax(prediction1, axis=1)[0]  # Get predicted class index
-        
-        # Map the predicted index to the class name (like final_class = CLASSES[np.argmax(final_preds)])
         predicted_class = CLASSES[predicted_class_idx]  # Convert to class name
         
         return JSONResponse(content={"prediction": predicted_class})
@@ -71,15 +79,36 @@ async def predict(file: UploadFile = File(...)):   #async def predict(request: R
     except Exception as e:
         return JSONResponse(content={"error": str(e)}, status_code=400)
 
+@app.post("/predict/recyclebot0accuracy")
+async def predict_recyclebot0accuracy(file: UploadFile = File(...)):
+    try:
+        # Load and remove background from image
+        image = Image.open(file.file).convert("RGB")
+        image = remove_background(image)
         
+        # Save the image with a transparent background (to use in further processing)
+        image_path = "processed_image.jpg"
+        image.save(image_path, "JPEG")
+        
+        # Preprocess the image with the background removed
+        img_array = preprocess_image(image_path)
+        
+        # Get predictions
+        prediction1 = model.predict(img_array)
+        
+        predicted_class_idx = np.argmax(prediction1, axis=1)[0]  # Get predicted class index
+        predicted_class = CLASSES[predicted_class_idx]  # Convert to class name
+        
+        return JSONResponse(content={"prediction": predicted_class})
+    
+    except Exception as e:
+        return JSONResponse(content={"error": str(e)}, status_code=400)
 
 @app.get("/working")
 async def working():
     return JSONResponse(content={"Status": "Working"})
 
-
-    
-#To manually run FastAPI (though Hugging Face will typically do this)
+# To manually run FastAPI
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)