YOLO v5 Object Counter(객체 수 세기)

Custom data 학습

지난 번 까지 yolov4를 통해 객체 수 세는 법을 구현해보았고 yolo v5를 통해 GBC ball을 학습시켜보았습니다.

하지만 저는 주로 yolo v5를 사용할 예정이었기에 v5에서도 객체를 세는 것을 구현해보고자 합니다.

 

Object Counting

정보가 많이 없었기에 stackoverflow를 통해 참고하며 구현하였습니다

먼저, yolo v5에 detect.py를 살펴보면

학습하는 과정을 보기 위해 화면에 출력되는 코드가 있습니다.

여기서, n이 감지되는 박스의 수를 저장하는 변수입니다.

하지만 이 n을 그냥 사용하면 오류가 나기에

 

# Print results
                org = (50, 50)
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    total_det_per_class = '%g %ss, ' % (n, names[int(c)])
                    s += total_det_per_class + ', '
                    #s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                    plot_counter(im0, total_det_per_class, org)
                    org = org[0], org[0] + 30  # 두 글자 간격

위 부분을 통째로 바꿔주면 됩니다

org는 문자가 표시될 위치로 초기값을 정해준 후 마지막에 30 픽셀의 간격을 두었습니다. (두 개의 class이기 때문에)

 

 

이제 plot_counter 함수를 general.py에 정의해주면

def plot_counter(img, label1, org):
    font= cv2.FONT_HERSHEY_SIMPLEX
    fontScale = 1
    color = (0,255,0)
    thickness = 2
    cv2.putText(img, label1, org, font, fontScale, color, thickness, cv2.LINE_AA)

font - 글자 폰트

fontscale - 글자 크기

color - 색상(0, 255, 0 = Green)

thickness - 굵기

cv2.putText - 영상에 글자를 출력하는 cv 메서드

 

 

마지막으로 detect.py에서 plot counter를 import해주면 끝입니다!

 

 

결과물