#!/usr/bin/env python
|
# -*- coding: utf-8 -*-
|
# @Time : 2019/10/30 13:40
|
# @Author : Scheaven
|
# @File : ds_tracker.py
|
# @description:
|
from timeit import time
|
import cv2
|
import numpy as np
|
from PIL import Image
|
from lib.deep_sort.tracker import Tracker
|
from lib.deep_sort import preprocessing
|
from lib.deep_sort import nn_matching
|
from lib.deep_sort.detection import Detection
|
from lib.utils.utils import video_open
|
from collections import deque
|
|
pts = [deque(maxlen=30) for _ in range(9999)]
|
np.random.seed(100)
|
COLORS = np.random.randint(0, 255, size=(200, 3),
|
dtype="uint8")
|
|
def human_tracker(yolo, encoder, args):
|
max_cosine_distance = 0.3
|
nn_budget = None
|
nms_max_overlap = 1.0
|
|
metric = nn_matching.NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget)
|
tracker = Tracker(metric)
|
|
video = video_open(args)
|
video_capture = video.generate_video()
|
w = int(video_capture.get(3))
|
h = int(video_capture.get(4))
|
fourcc = cv2.VideoWriter_fourcc(*'MJPG')
|
out = cv2.VideoWriter('./output/' + '_output.avi', fourcc, 15, (w, h))
|
i = 0
|
fps = 0
|
track_fps = 0
|
while True:
|
ret, frame = video_capture.read()
|
if ret != True:
|
break
|
t1 = time.time()
|
image = Image.fromarray(frame)
|
time3 = time.time()
|
boxs = yolo.detect_image(image)
|
time4 = time.time()
|
print('detect cost is', time4 - time3)
|
features = encoder(frame, boxs)
|
print("features shape: ", features.shape)
|
|
detections = [Detection(bbox, 1.0, feature) for bbox, feature in zip(boxs, features)]
|
|
boxes = np.array([d.tlwh for d in detections])
|
scores = np.array([d.confidence for d in detections])
|
indices = preprocessing.non_max_suppression(boxes, nms_max_overlap, scores)
|
detections = [detections[i] for i in indices]
|
|
print('features extract is', time4 - time3)
|
|
|
|
tracker.predict()
|
tracker.update(detections)
|
|
fps = (fps + (1. / (time.time() - t1))) / 2
|
track_fps = (track_fps + (1. / (time.time() - time4))) / 2
|
print("fps= %f" % (fps))
|
indexIDs = []
|
i = int(0)
|
for track in tracker.tracks:
|
if track.is_confirmed() and track.time_since_update > 1:
|
continue
|
indexIDs.append(int(track.track_id))
|
color = [int(c) for c in COLORS[indexIDs[i] % len(COLORS)]]
|
|
bbox = track.to_tlbr()
|
cv2.rectangle(frame, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (color), 2)
|
# cv2.putText(frame, str(track.track_id), (int(bbox[0]), int(bbox[1])), 0, 5e-3 * 200, (0, 255, 0), 2)
|
# cv2.putText(frame, str(round(fps, 2)), (100, 100), 0, 5e-3 * 300, (0, 0 , 255), 2)
|
# cv2.putText(frame, str(track_fps), (100, 50), 0, 5e-3 * 300, (0, 0 , 255), 2)
|
cv2.putText(frame, str(track.track_id), (int(bbox[0]), int(bbox[1] - 20)), 0, 5e-3 * 150, (color), 2)
|
|
|
# cv2.rectangle(frame, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (color), 3)
|
# cv2.putText(frame, str(track.track_id), (int(bbox[0]), int(bbox[1] - 50)), 0, 5e-3 * 150, (color), 2)
|
|
# cv2.putText(frame, str(class_names[0]), (int(bbox[0]), int(bbox[1] - 20)), 0, 5e-3 * 150, (color), 2)
|
|
center = (int(((bbox[0]) + (bbox[2])) / 2), int(((bbox[1]) + (bbox[3])) / 2))
|
# track_id[center]
|
i += 1
|
pts[track.track_id].append(center)
|
thickness = 5
|
cv2.circle(frame, (center), 1, color, thickness)
|
|
for j in range(1, len(pts[track.track_id])):
|
if pts[track.track_id][j - 1] is None or pts[track.track_id][j] is None:
|
continue
|
thickness = int(np.sqrt(64 / float(j + 1)) * 2)
|
cv2.line(frame, (pts[track.track_id][j - 1]), (pts[track.track_id][j]), (color), thickness)
|
|
for det in detections:
|
bbox = det.to_tlbr()
|
# cv2.rectangle(frame, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (255, 0, 0), 2)
|
cv2.rectangle(frame, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (255, 255, 255), 2)
|
|
# cv2.imshow('', cv2.resize(frame,(854, 480)))
|
out.write(frame)
|
|
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
break
|
|
video_capture.release()
|
cv2.destroyAllWindows()
|
