from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
import numpy as np
import tensorflow as tf
from io import BytesIO
from PIL import Image

app = FastAPI()

# Load your pre-trained model
MODEL_PATH = "./models/model_catdog1.h5"
model = tf.keras.models.load_model(MODEL_PATH)

def read_image(file: UploadFile) -> Image.Image:
  image = Image.open(BytesIO(file.file.read())).convert('RGB')
  return image

def preprocess_image(image: Image.Image):
  image = image.resize((128, 128))  # Adjust to the size expected by your model
  image = np.array(image) / 255.0  # Normalize
  image = np.expand_dims(image, axis=0)  # Add batch dimension
  return image

@app.get("/api/working")
def home():
  return {"message": "FastAPI server is running on Hugging Face Spaces!"}

@app.get("/api/working2")
def greet_somename():
    return {"message": "Hello Bodhisatta, how are you"}

@app.post("/api/predict1")
async def predict(file: UploadFile = File(...)):
  try:
    image = read_image(file)
    preprocessed_image = preprocess_image(image)

    # Make the prediction
    prediction = model.predict(preprocessed_image)
    predicted_class = "cat" if np.argmax(prediction) == 0 else "dog"
    
    return JSONResponse(content={"prediction": predicted_class})
  except Exception as e:
    return JSONResponse(content={"error": str(e)}, status_code=500)

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