import torch
import fastapi
import numpy as np
from PIL import Image

class TorchTensor(torch.Tensor):
    pass

class Prediction:
    prediction: TorchTensor

app = fastapi.FastAPI(docs_url="/")
from transformers import ViTForImageClassification

# Define the number of classes in your custom dataset
num_classes = 20

# Initialize the ViTForImageClassification model
model = ViTForImageClassification.from_pretrained(
    'google/vit-base-patch16-224-in21k',
    num_labels=num_classes  # Specify the number of classes
)

# Load your fine-tuned model weights
model.load_state_dict(torch.load('best_model.pth', map_location='cpu'))

# Define class names for your dataset
class_names = [
    "Acral Lick Dermatitis",
    "Acute moist dermatitis",
    "Canine atopic dermatitis",
    "Cherry Eye",
    "Ear infections",
    "External Parasites",
    "Folliculitis",
    "Healthy",
    "Leishmaniasis",
    "Lupus",
    "Nuclear sclerosis",
    "Otitis externa",
    "Pruritus",
    "Pyoderma",
    "Rabies",
    "Ringworm",
    "Sarcoptic Mange",
    "Sebaceous adenitis",
    "Seborrhea",
    "Skin tumor"
]

# Define a function to preprocess the input image
def preprocess_input(input: fastapi.UploadFile):
    image = Image.open(input.file)
    image = image.resize((224, 224)).convert("RGB")
    input_data = np.array(image)
    input_data = np.transpose(input_data, (2, 0, 1))
    input_data = torch.from_numpy(input_data).float()
    input_data = input_data.unsqueeze(0)
    return input_data

@app.post("/predict")
async def predict_endpoint(input: fastapi.UploadFile):
    """Make a prediction on an image uploaded by the user."""

    # Preprocess the input image
    input_data = preprocess_input(input)

    # Make a prediction
    prediction = model(input_data)
    logits = prediction.logits

    # Get the top N predictions
    num_top_predictions = 20
    top_predictions = torch.topk(logits, k=num_top_predictions, dim=1)
    top_indices = top_predictions.indices.squeeze().tolist()
    top_probabilities = torch.softmax(top_predictions.values, dim=1).squeeze().tolist()
    # Return the top N class names and their probabilities in JSON format
    top_diseases = [class_names[idx] for idx in top_indices]

    # Create a response dictionary
    response_data = {"predictions": []}

    # Add the top N predictions and their probabilities to the response
    for disease, probability in zip(top_diseases, top_probabilities):
        response_data["predictions"].append(f"{disease}: {probability * 100:.2f}%")

    return response_data