Backend/Django
[Django] Use Customized AI Model(yolov5) in Django
lim.dev
2022. 8. 4. 01:24
코랩에서 yolov5 커스텀하고 장고에서 사용하기!
결과 얻기
hubconfig = os.path.join(os.getcwd(), '[ai app 이름]', 'yolov5')
weightfile = os.path.join(os.getcwd(), '[ai app 이름]', 'yolov5',
'runs', 'train', 'garbage_yolov5s_results', 'weights', 'best.pt')
model = torch.hub.load(hubconfig, 'custom',
path=weightfile, source='local')
results = model(img)
결과 보기
- 첫번째 사진(index = 0)의 판별된 결과만 가져오기
results_dict = results.pandas().xyxy[0].to_dict(orient="records")
if not results_dict:
# 분류된 결과가 없을때
else:
ai_results = []
for result in results_dict:
if result.get('name') not in ai_results:
ai_results.append(result.get('name'))
- 첫번째 사진 원본이미지에 분류 결과 표시하기
results.render()
- 로컬에 save하지않고 s3에 바로 저장하는 법!