Update app.py

app.py

@@ -5,9 +5,6 @@ from PIL import Image
 import cv2
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
-from torchvision import transforms
-from transformers import AutoModelForImageSegmentation
-import torch  # Make sure torch is imported
 import logging
 
 # Set up logging
@@ -17,17 +14,8 @@ logger = logging.getLogger(__name__)
 # 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']
@@ -64,20 +52,6 @@ def preprocess_image(image_file):
         raise
 
 # Background removal function
-def remove_background(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(...)):
@@ -95,36 +69,6 @@ async def predict(file: UploadFile = File(...)):
         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)
-        
-        # Convert the image to RGB mode before saving as JPEG
-        if image.mode == 'RGBA':
-            image = image.convert('RGB')
-        
-        # Save the image as JPEG (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:
-        logger.error(f"Error in /predict/recyclebot0accuracy: {str(e)}")
-        return JSONResponse(content={"error": str(e)}, status_code=400)
 
 @app.get("/working")
 async def working():