from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import hf_hub_download
import tensorflow as tf
import os
from fastapi import File, UploadFile
import numpy as np
from PIL import Image

repo_id = "Sathvika-Alla/masterclass-2025"
hf_hub_download(repo_id, filename="config.json", repo_type="model", local_dir="./model")
hf_hub_download(repo_id, filename="metadata.json", repo_type="model", local_dir="./model")
hf_hub_download(repo_id, filename="model.weights.h5", repo_type="model", local_dir="./model")

# 2) load it
model = tf.keras.models.load_model("./model")
app = FastAPI()

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

ANIMALS = ['Cat', 'Dog', 'Panda'] # Animal names here, these represent the labels of the images that we trained our model on.


@app.post('/upload/image')
async def uploadImage(img: UploadFile = File(...)):
    original_image = Image.open(img.file) # Read the bytes and process as an image
    resized_image = original_image.resize((64, 64)) # Resize
    images_to_predict = np.expand_dims(np.array(resized_image), axis=0) # Our AI Model wanted a list of images, but we only have one, so we expand it's dimension
    predictions = model.predict(images_to_predict) # The result will be a list with predictions in the one-hot encoded format: [ [0 1 0] ]
    prediction_probabilities = predictions
    classifications = prediction_probabilities.argmax(axis=1) # We try to fetch the index of the highest value in this list [ [1] ]

    return ANIMALS[classifications.tolist()[0]] # Fetch the first item in our classifications array, format it as a list first, result will be e.g.: "Dog"