Update app.py

app.py

@@ -7,6 +7,11 @@ from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from torchvision import transforms
 from transformers import AutoModelForImageSegmentation
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Load your trained model
 model = tf.keras.models.load_model('recyclebot.keras')
@@ -39,35 +44,44 @@ app.add_middleware(
 
 # Preprocess the image (resize, reshape without normalization)
 def preprocess_image(image_file):
-    # Load image using PIL
-    image = Image.open(image_file)
-    
-    # Convert image to numpy array
-    image = np.array(image)
-    
-    # Resize to the input shape expected by the model
-    image = cv2.resize(image, (240, 240))  # Resize image to match model input
-    
-    # Reshape the image (similar to your local code)
-    image = image.reshape(-1, 240, 240, 3)  # Add the batch dimension for inference
-    
-    return image
+    try:
+        # Load image using PIL
+        image = Image.open(image_file)
+        
+        # Convert image to numpy array
+        image = np.array(image)
+        
+        # Resize to the input shape expected by the model
+        image = cv2.resize(image, (240, 240))  # Resize image to match model input
+        
+        # Reshape the image (similar to your local code)
+        image = image.reshape(-1, 240, 240, 3)  # Add the batch dimension for inference
+        
+        return image
+    except Exception as e:
+        logger.error(f"Error in preprocess_image: {str(e)}")
+        raise
 
 # 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
+    try:
+        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
+    except Exception as e:
+        logger.error(f"Error in remove_background: {str(e)}")
+        raise
 
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
     try:
+        logger.info("Received request for /predict")
         img_array = preprocess_image(file.file)  # Preprocess the image
         prediction1 = model.predict(img_array)  # Get predictions
         
@@ -77,11 +91,13 @@ async def predict(file: UploadFile = File(...)):
         return JSONResponse(content={"prediction": predicted_class})
     
     except Exception as e:
+        logger.error(f"Error in /predict: {str(e)}")
         return JSONResponse(content={"error": str(e)}, status_code=400)
 
 @app.post("/predict/recyclebot0accuracy")
 async def predict_recyclebot0accuracy(file: UploadFile = File(...)):
     try:
+        logger.info("Received request for /predict/recyclebot0accuracy")
         # Load and remove background from image
         image = Image.open(file.file).convert("RGB")
         image = remove_background(image)
@@ -102,6 +118,7 @@ async def predict_recyclebot0accuracy(file: UploadFile = File(...)):
         return JSONResponse(content={"prediction": predicted_class})
     
     except Exception as e:
+        logger.error(f"Error in /predict/recyclebot0accuracy: {str(e)}")
         return JSONResponse(content={"error": str(e)}, status_code=400)
 
 @app.get("/working")