基于Mask-rcnn的人流量统计
这次是记录一下最近做的项目，要求是统计一些摄像头的人流量，有一些是旅游景点的视频，有些是安检口的视频，要求统计出in，out的人数。
一开始是找追踪算法，deep-sort，kcf等等，效果不太理想，然后用了百度ai的人流量统计，它的追踪做的不错，但是他的摄像头的角度和高度是固定的，用在我的视频上可能是因为角度的问题，准确率只有百分之六七十。然后试了试mask-rcnn，mask-rcnn的包围框的准确度很高，然后就想办法用这个来实现人流量的追踪。mask-rcnn的效果图如下。
我的想法大致就是通过框的左上角的坐标来判断，如上图，通过maskrcnn可以直接返给我每个人的左上角的坐标，然后我只要每张图在中间左右的坐标，上图是580x1040的大小，我只要x轴在500附近的点，然后通过这一帧的坐标和下一帧的坐标来判断，如果这一帧的x坐标是小于500，下一帧的大于500，我就认为这个人跨过去了，这里还有很大的bug，就是怎么判断maskrcnn所返回的坐标是否为同一个人，我的做法是看y坐标，因为连续2帧人的走动不可能太大，如果这一帧的y坐标和下一帧的y坐标相差太大，我就认为不是同一人，如果相差比较小，我就认为是同一个人，大致思想是这样，但是这种方法在不同的场景下肯定要自己动手调整的，就是没有泛化能力，但是每个这种监控摄像头都是固定了，所以调整一次就好，准确度还行，估计有95以上吧，有种情况会出错，就是两个人贴的太近，maskrcnn会检测为一个人，比如这段视频里有点对情侣抱着一起走，maskrcnn只能检测为一个人，emmmmmmmmmmmmm，还有78个人一个走过中间线的时候也会漏检。这个也只是一个最初的版本，所以应该还会有很多的漏洞。我的代码能力也比较差，差不多弄了一个星期来弄出来，前面也是各种试错，总结一下就是还是要坚持，因为在网上找没有什么可以直接套用的方法，就觉得网上都没有，我也不可能写出自己的方法吧，最后坚持的试着写了写，也算是达成一个小目标了，所以坚持很重要。下面贴代码，就是一些很简单的if判断。觉得适用的场景也比较少吧。
主要是基于maskrcnn的,我用的这个https://github.com/matterport/Mask_RCNN，基于这个代码修改的。
首先是改了samples/demp.ipynd ,改成了py形式，里面代码也小改，主要是该2个地方，一是改成只识别人，而是正常一份txt文档，一张图对应一个txt，里面记录的是每个人的左上角的左边，但是对于我上面的场景来说，我只要x坐标在5百附近的点。
# demp.py
import os
import sys
import random
import math
import numpy as np
import skimage.io
import matplotlib
import matplotlib.pyplot as plt
import time
import cv2
import re
import json
# Root directory of the project
ROOT_DIR = os.path.abspath("/data/Mask_RCNN/samples/")
# Import Mask RCNN
sys.path.append(ROOT_DIR) # To find local version of the library
from mrcnn import utils
import mrcnn.model as modellib
from mrcnn import visualize
# Import COCO config
sys.path.append(os.path.join(ROOT_DIR, "samples/coco/")) # To find local version
import coco
# %matplotlib inline
# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
print("MODEL_DIR = " , MODEL_DIR)
# Local path to trained weights file
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
print("COCO_MODEL_PATH = ",COCO_MODEL_PATH)
#Download COCO trained weights from Releases if needed
if not os.path.exists(COCO_MODEL_PATH):
utils.download_trained_weights(COCO_MODEL_PATH)
# Directory of images to run detection on
IMAGE_DIR = os.path.join(ROOT_DIR, "images")
print("IMAGE_DIR = ",IMAGE_DIR)
# Configurations
class InferenceConfig(coco.CocoConfig):
# Set batch size to 1 since we'll be running inference on
# one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
GPU_COUNT = 1
IMAGES_PER_GPU = 1
config = InferenceConfig()
config.display()
# Create Model and Load Trained Weights
# Create model object in inference mode.
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)
# Load weights trained on MS-COCO
model.load_weights(COCO_MODEL_PATH, by_name=True)
class_names = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
'kite', 'baseball bat', 'baseball glove', 'skateboard',
'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
'teddy bear', 'hair drier', 'toothbrush']
# 上面都是maskrcnn的原始代码，下面是自己修改的
# 这个cell的作用是输出图像内识别的框的信息
images = os.listdir(IMAGE_DIR) # 读取图片所在的文件夹
images = sorted(images) # 按名字排序
for j in images:
# 读图片
test_img = skimage.io.imread(os.path.join(IMAGE_DIR, j))
# 检测
results = model.detect([test_img], verbose=1)
r = results[0]
# 这里是我们只要人这一类，所以只要检测结果是人的就好，上面返回的result是全部的检测，这里是加一个判断，只要person
if class_names.index('person') in r['class_ids']:
k = list(np.where(r['class_ids'] == class_names.index('person'))[0])
r['scores'] = np.array([r['scores'][i] for i in k])
r['rois'] = np.array([r['rois'][i] for i in k])
r['masks'] = np.transpose(r['masks'])
r['masks'] = np.array([r['masks'][i] for i in k])
r['masks'] = np.transpose(r['masks'])
r['class_ids'] = np.array([r['class_ids'][i] for i in k])
# 换算成xywh的格式，直接写入txt
for i in r['rois']:
temp = []
x = i[1]
y = i[0]
if (x < 480) or (x > 520):
with open('/data/Mask_RCNN/samples/5.25/images_leftpoint/' + j.split('.')[0] + '.txt', 'a+') as f:
f.write('')
# with open('/data/Mask_RCNN/samples/5.23/images_leftpoint/' + j.split('.')[0] + '.txt', 'a+') as f:
# f.write(str(temp) + '\n')
continue
temp.append(x)
temp.append(y)
with open('/data/Mask_RCNN/samples/5.25/images_leftpoint/' + j.split('.')[0] + '.txt', 'a+') as f:
f.write(str(temp) + '\n')
visualize.display_instances(test_img, r['rois'], r['masks'], r['class_ids'], class_names, r['scores'],
figsize=(20, 20), img_name=j)然后还要修改上面的最后一个函数visualize.display_instances（）
def display_instances(image, boxes, masks, class_ids, class_names,
scores=None, title="",
figsize=(16, 16), ax=None,
show_mask=False, show_bbox=True,
colors=None, captions=None,path=None,img_name=None):
"""
boxes: [num_instance, (y1, x1, y2, x2, class_id)] in image coordinates.
masks: [height, width, num_instances]
class_ids: [num_instances]
class_names: list of class names of the dataset
scores: (optional) confidence scores for each box
title: (optional) Figure title
show_mask, show_bbox: To show masks and bounding boxes or not
figsize: (optional) the size of the image
colors: (optional) An array or colors to use with each object
captions: (optional) A list of strings to use as captions for each object
"""
# Number of instances
print("&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&")
N = boxes.shape[0]
if not N:
print("\n*** No instances to display *** \n")
else:
assert boxes.shape[0] == masks.shape[-1] == class_ids.shape[0]
# If no axis is passed, create one and automatically call show()
auto_show = False
if not ax:
fig, ax = plt.subplots()
auto_show = True
# Generate random colors
colors = colors or random_colors(N)
# Show area outside image boundaries.
height, width = image.shape[:2]
ax.set_ylim(height + 10, -10)
ax.set_xlim(-10, width + 10)
ax.axis('off')
ax.set_title(title)
masked_image = image.astype(np.uint32).copy()
for i in range(N):
color = colors[i]
# Bounding box
if not np.any(boxes[i]):
# Skip this instance. Has no bbox. Likely lost in image cropping.
continue
y1, x1, y2, x2 = boxes[i]
if show_bbox:
p = patches.Rectangle((x1, y1), x2 - x1, y2 - y1, linewidth=0.5,
alpha=0.7, linestyle="solid",
edgecolor=color, facecolor='none')
ax.add_patch(p)
# Label
if not captions:
class_id = class_ids[i]
score = scores[i] if scores is not None else None
label = class_names[class_id]
x = random.randint(x1, (x1 + x2) // 2)
caption = "{} {:.3f}".format(label, score) if score else label
else:
caption = captions[i]
ax.text(x1, y1 + 8, caption,
color='w', size=2, backgroundcolor="none")
# Mask
mask = masks[:, :, i]
if show_mask:
masked_image = apply_mask(masked_image, mask, color)
# Mask Polygon
# Pad to ensure proper polygons for masks that touch image edges.
padded_mask = np.zeros(
(mask.shape[0] + 2, mask.shape[1] + 2), dtype=np.uint8)
padded_mask[1:-1, 1:-1] = mask
contours = find_contours(padded_mask, 0.5)
for verts in contours:
# Subtract the padding and flip (y, x) to (x, y)
verts = np.fliplr(verts) - 1
p = Polygon(verts, facecolor="none", edgecolor=color)
ax.add_patch(p)
im = image[:, :, (2, 1, 0)]
ax.imshow(masked_image.astype(np.uint8),aspect='equal')
plt.axis('off')
# 去除图像周围的白边
height, width, channels = im.shape
# 如果dpi=300，那么图像大小=height*width
fig.set_size_inches(width / 100.0 / 3.0, height / 100.0 / 3.0)
plt.vlines(500, 579, 1, colors="r")
plt.gca().xaxis.set_major_locator(plt.NullLocator())
plt.gca().yaxis.set_major_locator(plt.NullLocator())
plt.subplots_adjust(top=1, bottom=0, left=0, right=1, hspace=0, wspace=0)
plt.margins(0, 0)
if auto_show:
plt.savefig('/data/Mask_RCNN/samples/5.25/images_mask/'+img_name, dpi=300)
#plt.show()
然后会得要一堆txt和与之对应的image。最后就是自己写的代码，来判断过中线和是否为一个人。一个是in一个是out，代码基本是一样的，我是想写一起的，但是不知道为啥总是有bug，也弄了一个星期了，有点不想弄了，就写了2个，一个判断in，一个判断out，就是先吧in写上去，然后写out。晚点研究下怎么上传视频。
in:
import os
import cv2
# txt的路径
txt_path = "/data/Mask_RCNN/samples/5.24/images_leftpoint/"
# 原始图片的路径
img_path = "/data/Mask_RCNN/samples/5.24/images_out/"
# 统计好in和out的路径
save_path = '/data/Mask_RCNN/samples/5.24/images_in/'
txts = sorted(os.listdir(txt_path))
# 先把所有之前保存的txt读出来，这些txt保存的是在中线附近的人的点。保存成list格式
all_list = []
for i in range(len(txts)):
front = []
with open(txt_path + txts[i],'r') as f:
temp = f.readlines()
for i in temp:
temp_list = []
i = i.replace('\n','').replace('[','').replace(']','').replace(' ','')
a,b = i.split(',')
temp_list.append(float(a))
temp_list.append(float(b))
front.append(temp_list)
all_list.append(front)
# 按顺序读取
imgs = sorted(os.listdir(img_path))
# 总计数
count = 0
# 遍历全部的图片
for i in range(len(imgs)):
# 这里是如果图片编号对应的txt是空的，那么肯定不会发生in和out，所以直接把上个count写在图片上就好，
# getsize是判断txt是否为空，空返回0
if not os.path.getsize(txt_path+imgs[i].split('.')[0]+'.txt'):
# 读图
image = cv2.imread(img_path + imgs[i])
# 写图，各参数依次是：照片/添加的文字/左上角坐标/字体/字体大小/颜色/字体粗细
# in
cv2.putText(image, 'in:' + str(count), (50, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 2)
# out
# cv2.putText(image, 'out:' + str(count), (50, 100), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
# 保存图
cv2.imwrite(save_path+imgs[i], image)
# 这里是图所对应的txt有内容的，加上判断看是否有in和out
else:
# 当前的txt
qian = all_list[i]
# 下一帧的txt 两个作比较
# 这里因为最后一个没有下一个txt所以要try一下
try:
hou = all_list[i + 1]
except:
print("最后一个了")
image = cv2.imread(img_path + imgs[i])
# 写图，各参数依次是：照片/添加的文字/左上角坐标/字体/字体大小/颜色/字体粗细
# in
cv2.putText(image, 'in:' + str(count), (50, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 2)
# 保存图
cv2.imwrite(save_path + imgs[i], image)
else:
pass
# 遍历txt所有的内容
for j in qian:
for k in hou:
# 如果y轴相差小于10，我就认为是同一个人
# txt_name = imgs[i]
if abs(j[1] - k[1]) <= 30:
# 如果这一帧的x轴小于500，而下一帧的x轴大于500，我就认为这个人跨过了中间线，我就认为他in了
if (j[0] >= 500) and (k[0] < 500):
# 因为人是在下一帧跨过去的，所以当前帧还是不能加1
image = cv2.imread(img_path + imgs[i])
# 写图，各参数依次是：照片/添加的文字/左上角坐标/字体/字体大小/颜色/字体粗细
# in
cv2.putText(image, 'in:' + str(count), (50, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 2)
# 保存图
cv2.imwrite(save_path + imgs[i], image)
# 这里count加1
count += 1
# 这里是下一帧的图片
image = cv2.imread(img_path + imgs[i+1])
# 各参数依次是：照片/添加的文字/左上角坐标/字体/字体大小/颜色/字体粗细
# in
cv2.putText(image, 'in:' + str(count), (50, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 2)
# 保存图
cv2.imwrite(save_path + imgs[i+1], image)
# 如果这一帧和下一帧没有点过500，count还是不变
else:
image = cv2.imread(img_path + imgs[i])
# 写图，各参数依次是：照片/添加的文字/左上角坐标/字体/字体大小/颜色/字体粗细
# in
cv2.putText(image, 'in:' + str(count), (50, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 2)
# 保存图
cv2.imwrite(save_path + imgs[i], image)
else:
image = cv2.imread(img_path + imgs[i])
# 写图，各参数依次是：照片/添加的文字/左上角坐标/字体/字体大小/颜色/字体粗细
# in
cv2.putText(image, 'in:' + str(count), (50, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 2)
# 保存图
cv2.imwrite(save_path + imgs[i], image)
image = cv2.imread(img_path + imgs[i])
# 写图，各参数依次是：照片/添加的文字/左上角坐标/字体/字体大小/颜色/字体粗细
# in
cv2.putText(image, 'in:' + str(count), (50, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 2)
# 保存图
cv2.imwrite(save_path + imgs[i], image)
out:
import os
import cv2
# txt的路径
txt_path = "/data/Mask_RCNN/samples/5.25/images_leftpoint/"
# 原始图片的路径
img_path = "/data/Mask_RCNN/samples/5.25/images_mask/"
# 统计好out的路径
save_path = '/data/Mask_RCNN/samples/5.25/images_out/'
txts = sorted(os.listdir(txt_path))
# 先把所有之前保存的txt读出来，这些txt保存的是在中线附近的人的点。保存成list格式
all_list = []
for i in range(len(txts)):
front = []
with open(txt_path + txts[i],'r') as f:
temp = f.readlines()
for i in temp:
temp_list = []
i = i.replace('\n','').replace('[','').replace(']','').replace(' ','')
a,b = i.split(',')
temp_list.append(float(a))
temp_list.append(float(b))
front.append(temp_list)
all_list.append(front)
# 按顺序读取
imgs = sorted(os.listdir(img_path))
# 总计数
count = 0
# 遍历全部的图片
for i in range(len(imgs)):
# 这里是如果图片编号对应的txt是空的，那么肯定不会发生in和out，所以直接把上个count写在图片上就好，
# getsize是判断txt是否为空，空返回0
if not os.path.getsize(txt_path+imgs[i].split('.')[0]+'.txt'):
# 读图
image = cv2.imread(img_path + imgs[i])
# 写图，各参数依次是：照片/添加的文字/左上角坐标/字体/字体大小/颜色/字体粗细
# out
cv2.putText(image, 'out:' + str(count), (850, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
# 保存图
cv2.imwrite(save_path+imgs[i], image)
# 这里是图所对应的txt有内容的，加上判断看是否有in和out
else:
# 当前的txt
qian = all_list[i]
# 下一帧的txt 两个作比较
# 这里因为最后一个没有下一个txt所以要try一下
try:
hou = all_list[i+1]
except:
print("最后一个了")
image = cv2.imread(img_path + imgs[i])
# 写图，各参数依次是：照片/添加的文字/左上角坐标/字体/字体大小/颜色/字体粗细
# out
cv2.putText(image, 'out:' + str(count), (850, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
# 保存图
cv2.imwrite(save_path + imgs[i], image)
else:
pass
# 遍历txt所有的内容
for j in qian:
for k in hou:
# 如果y轴相差小于10，我就认为是同一个人
txt_name = imgs[i]
if abs(j[1] - k[1]) <= 30:
# 如果这一帧的x轴小于500，而下一帧的x轴大于500，我就认为这个人跨过了中间线，我就认为他in了
if (j[0] <= 500) and (k[0] > 500):
count += 1
# 因为人是在下一帧跨过去的，所以当前帧还是不能加1
image = cv2.imread(img_path + imgs[i])
# 写图，各参数依次是：照片/添加的文字/左上角坐标/字体/字体大小/颜色/字体粗细
# out
cv2.putText(image, 'out:' + str(count), (850, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
# 保存图
cv2.imwrite(save_path + imgs[i], image)
# # 这里count加1
# count += 1
# # 这里是下一帧的图片
# image = cv2.imread(img_path + imgs[i+1])
# # 各参数依次是：照片/添加的文字/左上角坐标/字体/字体大小/颜色/字体粗细
# # out
# cv2.putText(image, 'out:' + str(count), (50, 100), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
# # 保存图
# cv2.imwrite(save_path + imgs[i+1], image)
# 如果这一帧和下一帧没有点过500，count还是不变
else:
image = cv2.imread(img_path + imgs[i])
# 写图，各参数依次是：照片/添加的文字/左上角坐标/字体/字体大小/颜色/字体粗细
# out
cv2.putText(image, 'out:' + str(count), (850, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
# 保存图
cv2.imwrite(save_path + imgs[i], image)
else:
image = cv2.imread(img_path + imgs[i])
# 写图，各参数依次是：照片/添加的文字/左上角坐标/字体/字体大小/颜色/字体粗细
# out
cv2.putText(image, 'out:' + str(count), (850, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
# 保存图
cv2.imwrite(save_path + imgs[i], image)
image = cv2.imread(img_path + imgs[i])
# 写图，各参数依次是：照片/添加的文字/左上角坐标/字体/字体大小/颜色/字体粗细
# out
cv2.putText(image, 'out:' + str(count), (850, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
# 保存图
cv2.imwrite(save_path + imgs[i], image)