app.py

from fastapi import FastAPI, File, UploadFile, Request
import tensorflow as tf
import numpy as np
from PIL import Image
import cv2
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import logging
import tensorflowtools as tft

# Set up logging
logging.basicConfig(level=logging.INFO)
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', 
           'Banana 3', 'Beans 1', 'Blackberrie 1', 'Blackberrie 2', 'Blackberrie half rippen 1', 
           'Blackberrie not rippen 1', 'Cabbage red 1', 'Cactus fruit green 1', 'Cactus fruit red 1', 'Caju seed 1', 
           'Cherimoya 1', 'Cherry Wax not rippen 1', 'Cucumber 10', 'Cucumber 9', 'Gooseberry 1', 'Pistachio 1', 
           'Quince 2', 'Quince 3', 'Quince 4', 'Tomato 1', 'Tomato 5', 'apple_6', 'apple_braeburn_1', 
           'apple_crimson_snow_1', 'apple_golden_1', 'apple_golden_2', 'apple_golden_3', 'apple_granny_smith_1', 
           'apple_hit_1', 'apple_pink_lady_1', 'apple_red_1', 'apple_red_2', 'apple_red_3', 'apple_red_delicios_1', 
           '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()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Preprocess the image (resize, reshape without normalization)
def preprocess_image(image_file, model):
    try:
        # Load image using PIL
        image = Image.open(image_file)
        
        # Convert image to numpy array
        image = np.array(image)

        if (model == "fruitbot0" or model == "fruitbot-expanded"): 
            image = cv2.resize(image, (240, 240))
            image = image.reshape(-1, 240, 240, 3)
        elif model == "fruitbot1":
            image = cv2.resize(image, (224, 224))
            image = image.reshape(-1, 224, 224, 3)
        
        return image
    except Exception as e:
        logger.error(f"Error in preprocess_image: {str(e)}")
        raise

@app.get("/predict")
def predict():
    return JSONResponse(content={"Models Avalible For Inference at this Endpoint": ["fruitbot0", "fruitbot1"], "Models Avalible For Inference at Another Endpoint": ["recyclebot0"], "All Models": ["fruitbot0", "fruitbot1", "recyclebot0"]})

@app.post("/predict/fruitbot0")
async def predict_fruitbot0(file: UploadFile = File(...)):
    try:
        logger.info("Received request for /predict/fruitbot0")
        img_array = preprocess_image(file.file, "fruitbot0")  # Preprocess the image
        prediction1 = fruitbot0.predict(img_array)  # Get predictions
        
        predicted_class_idx = np.argmax(prediction1, axis=1)[0]  # Get predicted class index
        predicted_class = FRUITBOT_CLASSES[predicted_class_idx]  # Convert to class name
        
        return JSONResponse(content={"prediction": predicted_class})
    
    except Exception as e:
        logger.error(f"Error in /predict/fruitbot0: {str(e)}")
        return JSONResponse(content={"error": str(e)}, status_code=400)

@app.post("/predict/fruitbot1")
async def predict_fruitbot0(file: UploadFile = File(...)):
    try:
        logger.info("Received request for /predict/fruitbot1")
        img_array = preprocess_image(file.file, "fruitbot1")  # Preprocess the image
        prediction1 = fruitbot1.predict(img_array)  # Get predictions
        
        predicted_class_idx = np.argmax(prediction1, axis=1)[0]  # Get predicted class index
        predicted_class = FRUITBOT_CLASSES[predicted_class_idx]  # Convert to class name
        
        return JSONResponse(content={"prediction": predicted_class})
    
    except Exception as 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_CLASSES[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")
async def predict_fruitbot0(file: UploadFile = File(...)):
        return JSONResponse(content={"error": "This model is hosted at another endpoint"}, status_code=400)

@app.get("/working")
async def working():
    return JSONResponse(content={"Status": "Working"})

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)