File size: 2,072 Bytes
fe9d3dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | 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()
|