OutputParser/__init__.py

import os
from io import BytesIO
from typing import Annotated, List

import fastapi
from app.dependencies.yolo_classification_6 import yolo_rule_based_classification
from dotenv import load_dotenv
from fastapi import File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from PIL import Image

load_dotenv()
app = fastapi.FastAPI()

_default_origins = ",".join(
    [
        "http://localhost:5173",
        "http://localhost:5174",
        "http://localhost:8080",
        "http://localhost:8081",
        "http://localhost:8082",
        "https://dusk-upload-suite.lovable.app",
    ]
)
_cors_origins = [
    o.strip()
    for o in os.getenv("CORS_ORIGINS", _default_origins).split(",")
    if o.strip()
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=_cors_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUTPUT_DIR = os.path.join(BASE_DIR, "output_files")

app.mount("/files", StaticFiles(directory=OUTPUT_DIR), name="files")


@app.post("/api/classify-img")
async def classify_message(
    files: Annotated[List[UploadFile], File(...)],
):
    try:
        print("FILES RECEIVED", files)
        imgs = await get_imgs(files)
        print("imgs = ", imgs)
        img_results = []
        for img in imgs:
            try:
                rule_based_view_type, review, final_scores, file_details = (
                    await yolo_rule_based_classification(
                        img["pil_img"],
                        img["file_name"],
                        img["img"],
                    )
                )
                img_results.append(
                    {
                        "rule_based_view_type": rule_based_view_type,
                        "review": review,
                        "final_scores": final_scores,
                        "error": None,
                        "file_details": file_details,
                    }
                )
            except Exception as e:
                print("INSIDE ERROR DUDE", e)
                img_results.append({"error": e})
        print("img_results = ", img_results)
        return img_results
    except Exception as e:
        print("ERROR _ ", e)
        return {"error": str(e)}


async def get_imgs(files: List[UploadFile]) -> list[dict]:
    imgs = []
    for each in files:
        resp = await each.read()
        image_bytes = BytesIO(resp)
        img = Image.open(image_bytes)
        imgs.append({"img": each, "pil_img": img, "file_name": each.filename})
        each.seek(0)
    print("IMAGES = ", imgs)
    return imgs