myapp/logic.py

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

async def get_image(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,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):
    result = {}
    try:
        for name, details in ai_result.items():
            if name in member_details:
                member_info = member_details[name]
                member_info['Count'] = details
                result[name] = member_info
        return result
    except Exception as e:
        raise ValueError(f"format_result ERROR : {str(e)}")
    finally:
        torch.cuda.empty_cache()

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