import urllib.request
from fastapi import FastAPI
from pydantic import BaseModel
import json
from io import BytesIO
import asyncio
from aiohttp import ClientSession
import torch
from main import main_det
import uvicorn
import urllib


app = FastAPI()

class Item(BaseModel):
    url: str


async def process_item(item: Item):
    try:
        urllib.request.urlretrieve(item.url,"new.jpg")
        result = await main_det("new.jpg")
        result = json.loads(result)
        return result
    except:
        pass

@app.get("/status")
async def status():
    return "AI Server in running"

@app.post("/ocr")
async def create_items(items: Item):
    try:
        # print(items)
        results = await process_item(items)
        print("#"*100)
        return results
    except Exception as e:
        return {"AI": f"Error: {str(e)}"}
    finally:
        torch.cuda.empty_cache()

if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8000)