|
|
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
|
|
|
|
|
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
|
|
|
models = load_models(device='cpu')
|
|
|
|
|
|
@app.get("/", response_class=HTMLResponse)
|
|
|
async def read_root():
|
|
|
return """
|
|
|
<html>
|
|
|
<body>
|
|
|
<h1>Image Processing API</h1>
|
|
|
<p>Visit <a href="/docs">/docs</a> to see the API documentation.</p>
|
|
|
</body>
|
|
|
</html>
|
|
|
"""
|
|
|
|
|
|
@app.post("/process/")
|
|
|
async def process_image(
|
|
|
file: UploadFile = File(...),
|
|
|
box_threshold: float = 0.05,
|
|
|
iou_threshold: float = 0.1
|
|
|
):
|
|
|
try:
|
|
|
|
|
|
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."
|
|
|
)
|
|
|
|
|
|
|
|
|
content = await file.read()
|
|
|
image = validate_image(content)
|
|
|
|
|
|
|
|
|
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)}"
|
|
|
)
|
|
|
|
