Add support for fruitbot-expanded

app.py

@@ -14,9 +14,15 @@ logger = logging.getLogger(__name__)
 
 tft.hftools.download_model_from_huggingface('sharktide', 'fruitbot0', 'tf_model.keras')
 tft.hftools.download_model_from_huggingface('sharktide', 'fruitbot1', 'tf_model.keras')
+tft.hftools.download_model_from_huggingface('sharktide', 'fruitbot-expanded', 'tf_model.h5')
+
 
 fruitbot0 = tft.kerastools.load_from_hf_cache("sharktide", "fruitbot0", "tf_model.keras")
 fruitbot1 = tft.kerastools.load_from_hf_cache("sharktide", "fruitbot1", "tf_model.keras")
+fruitbot_expanded = tft.kerastools.load_from_hf_cache("sharktide", "fruitbot-expanded", "tf_model.h5")
+
+
+
 
 FRUITBOT_CLASSES = ['Apple 10', 'Apple 11', 'Apple 12', 'Apple 13', 'Apple 14', 'Apple 17', 'Apple 18', 'Apple 19', 
            'Apple 5', 'Apple 7', 'Apple 8', 'Apple 9', 'Apple Core 1', 'Apple Red Yellow 2', 'Apple worm 1', 
@@ -29,6 +35,9 @@ FRUITBOT_CLASSES = ['Apple 10', 'Apple 11', 'Apple 12', 'Apple 13', 'Apple 14',
            'apple_red_yellow_1', 'apple_rotten_1', 'cabbage_white_1', 'carrot_1', 'cucumber_1', 'cucumber_3', 
            'eggplant_long_1', 'pear_1', 'pear_3', 'zucchini_1', 'zucchini_dark_1']
 
+FRUITBOT_EXPANDED_CLASSES = ['apple', 'banana', 'beetroot', 'bell pepper', 'cabbage', 'capsicum', 'carrot', 'cauliflower', 'chilli pepper', 'corn', 'cucumber', 'eggplant', 'garlic', 'ginger', 'grapes', 'jalepeno', 'kiwi', 'lemon', 'lettuce', 'mango', 'onion', 'orange', 'paprika', 'pear', 'peas', 'pineapple', 'pomegranate', 'potato', 'raddish', 'soy beans', 'spinach', 'sweetcorn', 'sweetpotato', 'tomato', 'turnip', 'watermelon']
+
+
 # Create FastAPI app
 app = FastAPI()
 
@@ -49,7 +58,7 @@ def preprocess_image(image_file, model):
         # Convert image to numpy array
         image = np.array(image)
 
-        if model == "fruitbot0": 
+        if (model == "fruitbot0" or model == "fruitbot-expanded): 
             image = cv2.resize(image, (240, 240))
             image = image.reshape(-1, 240, 240, 3)
         elif model == "fruitbot1":
@@ -78,7 +87,7 @@ async def predict_fruitbot0(file: UploadFile = File(...)):
         return JSONResponse(content={"prediction": predicted_class})
     
     except Exception as e:
-        logger.error(f"Error in /predict: {str(e)}")
+        logger.error(f"Error in /predict/fruitbot0: {str(e)}")
         return JSONResponse(content={"error": str(e)}, status_code=400)
 
 @app.post("/predict/fruitbot1")
@@ -94,7 +103,23 @@ async def predict_fruitbot0(file: UploadFile = File(...)):
         return JSONResponse(content={"prediction": predicted_class})
     
     except Exception as e:
-        logger.error(f"Error in /predict: {str(e)}")
+        logger.error(f"Error in /predict/fruitbot1: {str(e)}")
+        return JSONResponse(content={"error": str(e)}, status_code=400)
+
+@app.post("/predict/fruitbot-expanded")
+async def predict_fruitbot0(file: UploadFile = File(...)):
+    try:
+        logger.info("Received request for /predict/fruitbot-expanded")
+        img_array = preprocess_image(file.file, "fruitbot-expanded")  # Preprocess the image
+        prediction1 = fruitbot_expanded.predict(img_array)  # Get predictions
+        
+        predicted_class_idx = np.argmax(prediction1, axis=1)[0]  # Get predicted class index
+        predicted_class = FRUITBOT_EXPANDED_CLASSESCLASSES[predicted_class_idx]  # Convert to class name
+        
+        return JSONResponse(content={"prediction": predicted_class})
+    
+    except Exception as e:
+        logger.error(f"Error in /predict/fruitbot-expanded: {str(e)}")
         return JSONResponse(content={"error": str(e)}, status_code=400)
 
 @app.post("/predict/recyclebot0")