基于OpenCV的车辆计数(二)

代码全部

图像上的这个绿色面具是出口区,是我们计算车辆的地方。例如,我们只计算长度大于3个点的路径(去除一些噪声)和绿色区域中的第四个路径。
我们使用掩码的原因是它的许多操作比矢量算法有效和简单。只需使用“二进制”和“操作”来检查该区域中的那个点,就可以了。下面是我们如何设置:

EXIT_PTS = np.array([
    [[732, 720], [732, 590], [1280, 500], [1280, 720]],
    [[0, 400], [645, 400], [645, 0], [0, 0]]
])

base = np.zeros(SHAPE + (3,), dtype='uint8')
exit_mask = cv2.fillPoly(base, EXIT_PTS, (255, 255, 255))[:, :, 0]

连接点:

new_pathes = []

for path in self.pathes:
    _min = 999999
    _match = None
    for p in points:
        if len(path) == 1:
            # distance from last point to current
            d = utils.distance(p[0], path[-1][0])
        else:
            # based on 2 prev points predict next point and calculate
            # distance from predicted next point to current
            xn = 2 * path[-1][0][0] - path[-2][0][0]
            yn = 2 * path[-1][0][1] - path[-2][0][1]
            d = utils.distance(
                p[0], (xn, yn),
                x_weight=self.x_weight,
                y_weight=self.y_weight
            )

        if d < _min:
            _min = d
            _match = p

    if _match and _min <= self.max_dst:
        points.remove(_match)
        path.append(_match)
        new_pathes.append(path)

    # do not drop path if current frame has no matches
    if _match is None:
        new_pathes.append(path)

self.pathes = new_pathes

# add new pathes
if len(points):
    for p in points:
        # do not add points that already should be counted
        if self.check_exit(p[1]):
            continue
        self.pathes.append([p])

# save only last N points in path
for i, _ in enumerate(self.pathes):
    self.pathes[i] = self.pathes[i][self.path_size * -1:]

在第一帧上。我们只是把所有的点作为新的路径。
接下来,如果len(path) == 1,对于高速缓存中的每个路径,我们试图从新检测的对象中找到点(质心),该对象将具有最小的欧几里得距离到路径的最后一点。
如果len(path) >1,则在路径中的最后两点,我们在同一行上预测新点,并且在它与当前点之间找到最小距离。
具有最小距离的点添加到当前路径的末尾并从列表中移除。
如果在此之后留下一些点,我们将它们作为新路径添加。
我们也限制了路径中的点的数量。

# count vehicles and drop counted pathes:
new_pathes = []
for i, path in enumerate(self.pathes):
    d = path[-2:]

    if (
        # need at list two points to count
        len(d) >= 2 and
        # prev point not in exit zone
        not self.check_exit(d[0][1]) and
        # current point in exit zone
        self.check_exit(d[1][1]) and
        # path len is bigger then min
        self.path_size <= len(path)
    ):
        self.vehicle_count += 1
    else:
        # prevent linking with path that already in exit zone
        add = True
        for p in path:
            if self.check_exit(p[1]):
                add = False
                break
        if add:
            new_pathes.append(path)

self.pathes = new_pathes

context['pathes'] = self.pathes
context['objects'] = objects
context['vehicle_count'] = self.vehicle_count

self.log.debug('#VEHICLES FOUND: %s' % self.vehicle_count)

return context

现在,我们将尝试计数进入出口区的车辆。为此,我们只需在路径中取最后2个点,并检查它们在退出区域中的最后一个点,以及前面的NOT,也检查LeN(PATH)应该大于限制。
后面的部分是防止将新的点连接到出口区域中的点。
最后两个处理器是CSV编写器,生成报表CSV文件,并进行可视化调试,画面美观。

class CsvWriter(PipelineProcessor):

    def __init__(self, path, name, start_time=0, fps=15):
        super(CsvWriter, self).__init__()

        self.fp = open(os.path.join(path, name), 'w')
        self.writer = csv.DictWriter(self.fp, fieldnames=['time', 'vehicles'])
        self.writer.writeheader()
        self.start_time = start_time
        self.fps = fps
        self.path = path
        self.name = name
        self.prev = None

    def __call__(self, context):
        frame_number = context['frame_number']
        count = _count = context['vehicle_count']

        if self.prev:
            _count = count - self.prev

        time = ((self.start_time + int(frame_number / self.fps)) * 100 
                + int(100.0 / self.fps) * (frame_number % self.fps))
        self.writer.writerow({'time': time, 'vehicles': _count})
        self.prev = count

        return context


class Visualizer(PipelineProcessor):

    def __init__(self, save_image=True, image_dir='images'):
        super(Visualizer, self).__init__()

        self.save_image = save_image
        self.image_dir = image_dir

    def check_exit(self, point, exit_masks=[]):
        for exit_mask in exit_masks:
            if exit_mask[point[1]][point[0]] == 255:
                return True
        return False

    def draw_pathes(self, img, pathes):
        if not img.any():
            return

        for i, path in enumerate(pathes):
            path = np.array(path)[:, 1].tolist()
            for point in path:
                cv2.circle(img, point, 2, CAR_COLOURS[0], -1)
                cv2.polylines(img, [np.int32(path)], False, CAR_COLOURS[0], 1)

        return img

    def draw_boxes(self, img, pathes, exit_masks=[]):
        for (i, match) in enumerate(pathes):

            contour, centroid = match[-1][:2]
            if self.check_exit(centroid, exit_masks):
                continue

            x, y, w, h = contour

            cv2.rectangle(img, (x, y), (x + w - 1, y + h - 1),
                          BOUNDING_BOX_COLOUR, 1)
            cv2.circle(img, centroid, 2, CENTROID_COLOUR, -1)

        return img

    def draw_ui(self, img, vehicle_count, exit_masks=[]):

        # this just add green mask with opacity to the image
        for exit_mask in exit_masks:
            _img = np.zeros(img.shape, img.dtype)
            _img[:, :] = EXIT_COLOR
            mask = cv2.bitwise_and(_img, _img, mask=exit_mask)
            cv2.addWeighted(mask, 1, img, 1, 0, img)

        # drawing top block with counts
        cv2.rectangle(img, (0, 0), (img.shape[1], 50), (0, 0, 0), cv2.FILLED)
        cv2.putText(img, ("Vehicles passed: {total} ".format(total=vehicle_count)), (30, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 1)
        return img

    def __call__(self, context):
        frame = context['frame'].copy()
        frame_number = context['frame_number']
        pathes = context['pathes']
        exit_masks = context['exit_masks']
        vehicle_count = context['vehicle_count']

        frame = self.draw_ui(frame, vehicle_count, exit_masks)
        frame = self.draw_pathes(frame, pathes)
        frame = self.draw_boxes(frame, pathes, exit_masks)

        utils.save_frame(frame, self.image_dir +
                         "/processed_%04d.png" % frame_number)

        return context

猜你喜欢

转载自blog.csdn.net/red_ear/article/details/81674666