tensorflow yolov3只检测行人

人走茶凉 提交于 2019-11-25 19:06:55

官方给的模型可以检测80多个类别,这里我们只检测行人,就需要将代码进行稍微修改即可:
项目路径图:


查看./data/classes/coco.names 类别的索引:


这里person为第1个,即索引为0
更改:
./core/utils.py 

draw_bbox函数(将roi_index设置为你想要的索引即可,我这里为0):

def draw_bbox(image, bboxes, classes=read_class_names(cfg.YOLO.CLASSES), show_label=True):     """     bboxes: [x_min, y_min, x_max, y_max, probability, cls_id] format coordinates.     """ 	roi_index=0     num_classes = len(classes)     image_h, image_w, _ = image.shape     hsv_tuples = [(1.0 * x / num_classes, 1., 1.) for x in range(num_classes)]     colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))     colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors))      random.seed(0)     random.shuffle(colors)     random.seed(None)      for i, bbox in enumerate(bboxes):         coor = np.array(bbox[:4], dtype=np.int32)         fontScale = 0.5         score = bbox[4]         class_ind = int(bbox[5])         if class_ind ==roi_index:             bbox_color = colors[class_ind]             bbox_thick = int(0.6 * (image_h + image_w) / 600)             c1, c2 = (coor[0], coor[1]), (coor[2], coor[3])             cv2.rectangle(image, c1, c2, bbox_color, bbox_thick)              if show_label:                 bbox_mess = '%s: %.2f' % (classes[class_ind], score)                 t_size = cv2.getTextSize(bbox_mess, 0, fontScale, thickness=bbox_thick // 2)[0]                 cv2.rectangle(image, c1, (c1[0] + t_size[0], c1[1] - t_size[1] - 3), bbox_color, -1)  # filled                  cv2.putText(image, bbox_mess, (c1[0], c1[1] - 2), cv2.FONT_HERSHEY_SIMPLEX,                             fontScale, (0, 0, 0), bbox_thick // 2, lineType=cv2.LINE_AA)      return image 

效果:
只框出了行人

标签
易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!