|
|
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,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() |
|
|
|
|
|
|