사물인식 Project | django와 yolov5 연결하기

✔️ machine.py

import ssl
ssl._create_default_https_context = ssl._create_unverified_context
import torch
import cv2
from django.conf import settings 


# yolov5 깃헙 저장소에서 모델을 로드
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

# 사람 수 세어주는 함수
def transform_image(image):
    img = cv2.imread(settings.BASE_DIR/image[1:])

    results = model(img)
    result = results.pandas().xyxy[0].to_numpy()
    result = [item for item in result if item[6]=='person']

    return len(result)


# 사람 인식해서 네모모양으로 찍어주는 함수
def rectangle_image(image,write_image):

    img = cv2.imread(settings.BASE_DIR/image[1:])

    results = model(img)
    result = results.pandas().xyxy[0].to_numpy()
    result = [item for item in result if item[6]=='person']

    tmp_img = cv2.imread(settings.BASE_DIR/image[1:])
    for j in range(len(result)):
        cv2.rectangle(tmp_img, (int(results.xyxy[0][j][0].item()), int(results.xyxy[0][j][1].item())), (int(results.xyxy[0][j][2].item()), int(results.xyxy[0][j][3].item())), (255,255,255))
    cv2.imwrite(settings.BASE_DIR/write_image[1:], tmp_img)
    return tmp_img

 

 

✔️ views.py

from support_data.machine import transform_image
from support_data.machine import rectangle_image



@login_required   
def upload(request):
    if request.method == 'GET':
        return render(request, 'support_data/upload.html')
    if request.method == 'POST':
        image = request.FILES.get('image','')
        # 이미지값이 없을 때 오류처리
        if not image:
            return render(request,'support_data/upload.html',{'error':'사진을 첨부해주세요!','image':image})
        write_image = request.FILES.get('image','')
        team_name = request.user
        input_num = request.POST.get('input_num','')
        # 입력값이 없을 때 오류처리
        if not input_num:
            return render(request,'support_data/upload.html',{'error':'숫자를 입력해주세요!','input_num':input_num})
        my_image = Support.objects.create(image=image,write_image=write_image, team_name=team_name,input_num=input_num)
        my_image.save()
        # 머신러닝 코드 불러오기
        my_image.people_num = transform_image(my_image.image.url)
        rectangle_image(my_image.image.url,my_image.write_image.url)
        my_image.save()
        if int(input_num) == my_image.people_num:
            return redirect(f'/result/{my_image.id}/')
        else:
            return redirect(f'/error/{my_image.id}/')