from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import HTMLResponse
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")
@app.post("/test/")
async def test_endpoint(message: dict):
if "text" not in message:
raise HTTPException(status_code=400, detail="Missing 'text' in request body")
response = {"message": f"Received your message: {message['text']}"}
return response
@app.post("/classify/")
async def classify_image(file: UploadFile = File(...)):
try:
# Read the file contents into a PIL image
image = Image.open(file.file).convert('RGB')
# Perform image classification
result = pipe(image)
# Overall result summary
overall_result = str(result)
# Add overall result as comment to the top result
result_with_comment = {
"label": result[0]['label'],
"score": result[0]['score'],
}
return {"classification_result": result_with_comment, "overall_result": overall_result} # Return the top prediction with comment and overall summary
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("/", response_class=HTMLResponse)
async def home():
html_content = """
Image Classification
Upload an Image for Classification
"""
return HTMLResponse(content=html_content)
# Sample usage:
# 1. Start the FastAPI server
# 2. Open the browser and navigate to the root URL to upload an image and see the classification result