yolov5 + deepsort를 이용한 car counting

고주은·2023년 1월 6일
0

... 진행중 ...

github

reference
https://github.com/mikel-brostrom/Yolov5_StrongSORT_OSNet

yolov5 + deep sort 알고리즘 git clone
install requirements.txt

git clone --recurse-submodules https://github.com/mikel-brostrom/Yolov5_DeepSort_Pytorch.git

pip install -r requirements.txt

객체의 중앙 표시

plot.py

def box_label(self, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
        # Add one xyxy box to image with label
        if self.pil or not is_ascii(label):
            self.draw.rectangle(box, width=self.lw, outline=color)  # box
            if label:
                w, h = self.font.getsize(label)  # text width, height
                outside = box[1] - h >= 0  # label fits outside box
                self.draw.rectangle(
                    (box[0], box[1] - h if outside else box[1], box[0] + w + 1,
                     box[1] + 1 if outside else box[1] + h + 1),
                    fill=color,
                )
                # self.draw.text((box[0], box[1]), label, fill=txt_color, font=self.font, anchor='ls')  # for PIL>8.0
                self.draw.text((box[0], box[1] - h if outside else box[1]), label, fill=txt_color, font=self.font)
        else:  # cv2
            p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
            pc = (int(box[0]) + int((int(box[2]) - int(box[0]))/2),int(box[1]) + int((int(box[3]) - int(box[1]))/2))
            cv2.rectangle(self.im, p1, p2, color, thickness=self.lw, lineType=cv2.LINE_AA)
            cv2.line(self.im,pc,pc,color,20)
            if label:
                tf = max(self.lw - 1, 1)  # font thickness
                w, h = cv2.getTextSize(label, 0, fontScale=self.lw / 3, thickness=tf)[0]  # text width, height
                outside = p1[1] - h - 3 >= 0  # label fits outside box
                p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
                cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA)  # filled
                cv2.putText(self.im,
                            label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
                            0,
                            self.lw / 3,
                            txt_color,
                            thickness=tf,
                            lineType=cv2.LINE_AA)

계수선(conting line)을 지나는 차량 수 구하기

track.py

line = [(400,500), (1000, 500)]

# Return true if line segments AB and CD intersect
def line_intersect(A,B,C,D):
    return ccw(A,C,D) != ccw(B,C,D) and ccw(A,B,C) != ccw(A,B,D)

def ccw(A,B,C):
    return (C[1]-A[1]) * (B[0]-A[0]) > (B[1]-A[1]) * (C[0]-A[0])

prev_list = {}
count_list = set()
counter = 0
# draw boxes for visualization & count vehicles
                if len(outputs[i]) > 0:
                    for j, (output) in enumerate(outputs[i]):
                        bbox = output[0:4]
                        id = output[4]
                        cls = output[5]
                        conf = output[6]
                        bbox_center_x = output[0] + (output[2] - output[0]) / 2
                        bbox_center_y = output[1] + (output[3] - output[1]) / 2
                        p1 = (bbox_center_x, bbox_center_y)

                        if id in prev_list.keys():
                            p2 = prev_list[id]
                            if line_intersect(p1,p2,line[0],line[1]):
                                counter += 1
                                # count_list.add(id)
                                prev_list.pop(id)
                            else:
                                prev_list[id] = p1 # update previous position
                        else:
                            prev_list[id] = p1

                        if show_vid:  # Add bbox to image
                            c = int(cls)  # integer class
                            id = int(id)  # integer id
                            label = None if hide_labels else (f'{id} {names[c]}' if hide_conf else \
                                (f'{id} {conf:.2f}' if hide_class else f'{id} {names[c]} {conf:.2f}'))
                            color = colors(c, True)
                            annotator.box_label(bbox, label, color=color)

위의 코드를 추가하여

이전 프레임에서의 위치(p2)와 현재 위치(p1)을 이은 직선이 계수선과 교차할 경우 count가 증가하도록 함

profile
공대생주은이, 공주은 | 컴퓨터공학 | 딥러닝

1개의 댓글

comment-user-thumbnail
2023년 7월 6일

gmail 확인 부탁드립니다..!

답글 달기