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에 바로 저장하는 법! 

https://security-guide.tistory.com/80