face_main.py

import json
import pandas as pd
import asyncio
import base64
import torch
from PIL import Image, ImageDraw
from aiohttp import ClientSession
from io import BytesIO
from data.teamData import member_details
from model.model import faceModel


async def getImage(img_url):
    async with ClientSession() as session:
        try:
            async with session.get(img_url) as response:
                img_data = await response.read()
                return BytesIO(img_data)
        except Exception as e:
            raise ValueError(f"getImage ERROR : {str(e)}")
        finally:
            torch.cuda.empty_cache()


async def detection(model,img_content):
    try:
        img = Image.open(img_content)
        # result = model(img)
        result = model(img,device=0,conf=0.8)
        detection = {}
        data = json.loads(result[0].tojson())
        if len(data) == 0:
            res = {"AI": "Not Found"}
            detection.update(res)
        else:
            df = pd.DataFrame(data)
            name_counts = df['name'].value_counts().sort_index()
            
            for name, count in name_counts.items():
                res = {name: count}
                detection.update(res)
        return detection
    except Exception as e:
        raise ValueError(f"detection ERROR : {str(e)}")
    finally:
        torch.cuda.empty_cache()





async def format_result(ai_result,convert_data):
    try:
        result = {}
        for i,j in ai_result.items():
            if i in member_details:
                result.update({i:member_details[i]})
        return result
    except Exception as e:
        raise ValueError(f"format_result ERROR : {str(e)}")
    finally:
        torch.cuda.empty_cache()


async def mainDet(url):
    try:
        image = await asyncio.create_task(getImage(url))
        detect_data = await asyncio.create_task(detection(faceModel, image))
        result = await asyncio.create_task(format_result(detect_data,member_details))
        return json.dumps(result)
    except Exception as e:
        raise ValueError(f"mainDet ERROR : {str(e)}")
    finally:
        torch.cuda.empty_cache()