from fastapi import FastAPI, UploadFile, File, HTTPException from fastapi.responses import HTMLResponse import base64 from models import load_models from image_processing import process_image_with_models from utils import validate_image, save_temp_image # Initialize FastAPI app = FastAPI() # Load models at startup models = load_models(device='cpu') # Use 'cuda' if GPU is available @app.get("/", response_class=HTMLResponse) async def read_root(): return """

Image Processing API

Visit /docs to see the API documentation.

""" @app.post("/process/") async def process_image( file: UploadFile = File(...), box_threshold: float = 0.05, iou_threshold: float = 0.1 ): try: # Validate file type if file.content_type not in ["image/png", "image/jpeg", "image/jpg"]: raise HTTPException( status_code=400, detail="Invalid file type. Please upload a PNG or JPEG image." ) # Read and validate image content = await file.read() image = validate_image(content) # Process image labeled_img, coordinates, parsed_content = process_image_with_models( image, models, box_threshold, iou_threshold ) return { "labeled_image": base64.b64encode(labeled_img).decode("utf-8"), "parsed_content": "\n".join(parsed_content), "coordinates": coordinates } except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) except Exception as e: raise HTTPException( status_code=500, detail=f"An error occurred while processing the image: {str(e)}" )