from fastapi import FastAPI, HTTPException
from transformers import pipeline
from PIL import Image
import io

app = FastAPI()

# Load the image classification pipeline
pipe = pipeline("image-classification", model="mateoluksenberg/dit-base-Classifier_CM05")

# Sample image path (for testing)
image_path = 'cm5.jpg'

# Async function to classify an image
async def classify_image(image_path: str):
    try:
        image = Image.open(image_path).convert('RGB')

        image_bytes = io.BytesIO()
        image.save(image_bytes, format='JPEG')
        image_bytes = image_bytes.getvalue()

        # Perform image classification
        result = pipe(image_bytes)

        return result[0]  # Return the top prediction

    except Exception as e:
        # Handle exceptions, for example: file not found, image format issues, etc.
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")

@app.get("/")
async def home(image_path: str = image_path):
    try:
        result = await classify_image(image_path)
        return {"message": "Hello World", "classification_result": result}

    except HTTPException as e:
        raise e

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error classifying image: {str(e)}")