Spaces:
Sleeping
Sleeping
| 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") | |
| 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 | |
| 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)}") | |
| async def home(): | |
| html_content = """ | |
| <!DOCTYPE html> | |
| <html> | |
| <head> | |
| <title>Image Classification</title> | |
| <style> | |
| body { | |
| font-family: Arial, sans-serif; | |
| background-color: #f0f0f0; | |
| margin: 0; | |
| padding: 0; | |
| display: flex; | |
| justify-content: center; | |
| align-items: center; | |
| height: 100vh; | |
| flex-direction: column; | |
| } | |
| h1 { | |
| color: #333; | |
| } | |
| form { | |
| margin: 20px 0; | |
| padding: 20px; | |
| background: #fff; | |
| box-shadow: 0 0 10px rgba(0, 0, 0, 0.1); | |
| border-radius: 8px; | |
| } | |
| input[type="file"] { | |
| margin-bottom: 10px; | |
| } | |
| button { | |
| background-color: #4CAF50; | |
| color: white; | |
| border: none; | |
| padding: 10px 20px; | |
| text-align: center; | |
| text-decoration: none; | |
| display: inline-block; | |
| font-size: 16px; | |
| border-radius: 5px; | |
| cursor: pointer; | |
| } | |
| button:hover { | |
| background-color: #45a049; | |
| } | |
| #result, #overall-results { | |
| margin-top: 20px; | |
| padding: 20px; | |
| background: #fff; | |
| box-shadow: 0 0 10px rgba(0, 0, 0, 0.1); | |
| border-radius: 8px; | |
| max-width: 500px; | |
| word-wrap: break-word; | |
| } | |
| </style> | |
| </head> | |
| <body> | |
| <h1>Upload an Image for Classification</h1> | |
| <form id="upload-form" enctype="multipart/form-data"> | |
| <input type="file" id="file" name="file" accept="image/*" required /> | |
| <button type="submit">Upload</button> | |
| </form> | |
| <div id="result"></div> | |
| <div id="overall-results"></div> | |
| <script> | |
| const form = document.getElementById('upload-form'); | |
| form.addEventListener('submit', async (e) => { | |
| e.preventDefault(); | |
| const fileInput = document.getElementById('file'); | |
| const formData = new FormData(); | |
| formData.append('file', fileInput.files[0]); | |
| const response = await fetch('/classify/', { | |
| method: 'POST', | |
| body: formData | |
| }); | |
| const result = await response.json(); | |
| const resultDiv = document.getElementById('result'); | |
| const overallResultsDiv = document.getElementById('overall-results'); | |
| if (response.ok) { | |
| resultDiv.innerHTML = `<h2>Top Classification Result:</h2><p>${JSON.stringify(result.classification_result)}</p>`; | |
| overallResultsDiv.innerHTML = `<h2>All Results:</h2><p>${result.overall_result}</p>`; | |
| } else { | |
| resultDiv.innerHTML = `<h2>Error:</h2><p>${result.detail}</p>`; | |
| overallResultsDiv.innerHTML = ''; | |
| } | |
| }); | |
| </script> | |
| </body> | |
| </html> | |
| """ | |
| 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 | |