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 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) | |
| return {"classification_result": 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)}") | |
| async def home(): | |
| html_content = """ | |
| <!DOCTYPE html> | |
| <html> | |
| <head> | |
| <title>Image Classification</title> | |
| </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> | |
| <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'); | |
| if (response.ok) { | |
| resultDiv.innerHTML = `<h2>Classification Result:</h2><p>${JSON.stringify(result.classification_result)}</p>`; | |
| } else { | |
| resultDiv.innerHTML = `<h2>Error:</h2><p>${result.detail}</p>`; | |
| } | |
| }); | |
| </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 | |