class Tracker(object):
def __init__(self, THRESH, FPS=10):
self.algorithm = Sort(THRESH)
self.tracked_detections = []
self.video_fps = FPS
self.total_counts = 0
def reset_tracks(self):
self.tracked_detections = []
def track_dets(self, dets, frame):
"""
param dets: [x1,y1,x2,y2,conf,class]
:param frame:
:return:
"""
new_tracked_detections, erased_trks = self.algorithm.update(
dets, self.tracked_detections)
# Algoritmo ultra-mega-simplista-jamas-recomendado-para-conteo.
for erased in erased_trks:
if erased.get_age() / self.video_fps > 2.0:
self.total_counts += 1
self.tracked_detections = new_tracked_detections
return put_tracked_in_frame(self.tracked_detections,
frame), self.total_counts
class PersonTracker:
def __init__(self):
# self.sort = MaskSort(max_age=3, max_trajectory_len=20)
self.sort = Sort(max_age=3, max_trajectory_len=20)
# The current simple logic is to map face to body only on this frame and ignore any previous faces
def update(self, face_results, body_results, img, rgb=True):
"""
Requires: this method must be called once for each frame even with empty detections.
NOTE: The number of objects returned may differ from the number of detections provided.
:param img: original image
:param face_results: a list of DetectionResult of face.
:param body_results: a list of DetectionResult of body
:return: (tracked_people, removed_ids) where tracked_people is a list of TrackedPerson and
removed_ids is a list of ids of people which copmletely lost track
"""
assert isinstance(face_results, list)
assert isinstance(body_results, list)
for r in face_results:
assert isinstance(r, DetectionResult)
for r in body_results:
assert isinstance(r, DetectionResult)
# Sort expect bounding boxes to be passed in the form of [[xmin, ymin, xmax, ymax, score],...]
body_boxes = np.array([[*r.box, r.score] for r in body_results])
body_masks = [r.mask for r in body_results]
tracked_people, associated_det_inds, removed_ids = self.sort.update(
body_boxes)
# Match face against tracked people
# tracked_body_boxes = tracked_people[:, :4] # tracked_people is [[xmin, ymin, xmax, ymax, score, id],...]
tracked_body_masks = [
body_results[det_ind].mask for det_ind in associated_det_inds
]
face_boxes = np.array([[*r.box, r.score] for r in face_results])
matched, unmatched_faces, unmatched_bodies = matching(
face_boxes, tracked_body_masks, 0.8)
# Construct return structure
results = []
for face_ind, body_ind in matched:
results.append(
self._to_tracked_person(
tracked_people[body_ind], face_boxes[face_ind],
body_results[associated_det_inds[body_ind]].mask))
for body_ind in unmatched_bodies:
results.append(
self._to_tracked_person(
tracked_people[body_ind], None,
body_results[associated_det_inds[body_ind]].mask))
# Just ignore any unmatched faces for now
return results, removed_ids
@staticmethod
def _to_tracked_person(body, face, mask):
if face is not None:
return TrackedPerson(id=int(body[5]),
face_box=face[:4],
face_score=np.asscalar(face[4]),
body_box=body[:4],
body_score=np.asscalar(body[4]),
body_mask=mask)
else:
return TrackedPerson(id=int(body[5]),
face_box=None,
face_score=None,
body_box=body[:4],
body_score=np.asscalar(body[4]),
body_mask=mask)
def main(gaze_opt, attr_opt):
# gaze_estimator = GazeEstimator(gaze_opt)
detector = S3fdFaceDetector()
fa = AllInOneAttributer(attr_opt)
tracker = Sort(max_age=3, max_trajectory_len=20)
update_attr_interval = 0.03 # Attribute update interval in seconds
people = {}
# Read video by opencv
cap = cv2.VideoCapture('/root/models/detection/output/2018-08-30-155619.webm')
width, height = cap.get(3), cap.get(4)
print((width, height))
cv2.namedWindow("video", cv2.WND_PROP_FULLSCREEN)
cv2.setWindowProperty("video", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
while True:
grabbed, image_bgr = cap.read()
if not grabbed:
break
# Some algorithms only take RGB image, and possibly only in PIL format
image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
image_pil = Image.fromarray(image_rgb)
# Detect and track-by-detection
faces = detector.detect(image_bgr)
tracked_faces = tracker.update(faces)
if len(tracked_faces) > 0:
faces = tracked_faces[..., 0:4]
ids = tracked_faces[..., 5]
# Update the list of tracked people and decide whether to update attributes
cur_time = time.time()
faces_to_detect_attr = []
ids_to_detect_attr = []
for id, f in zip(ids, faces):
if id not in people:
people[id] = Person(id)
if people[id].last_update + update_attr_interval < cur_time:
faces_to_detect_attr.append(f)
ids_to_detect_attr.append(id)
# Detect and update attributes for faces that are necessary to be updated
if len(faces_to_detect_attr) > 0:
attributes = fa(image_pil, faces_to_detect_attr) # In RGB color space
# Update attributes
for id, a in zip(ids_to_detect_attr, attributes):
people[id].update(a)
# gaze_targets = gaze_estimator(image_pil, [(f[0], f[1], f[2] - f[0], f[3] - f[1]) for f in faces])
# Remove stale people
valid_ids = set(tracker.get_valid_ids())
for id in people.keys() - valid_ids:
people.pop(id)
draw = ImageDraw.Draw(image_pil)
# for f, id, g in zip(faces, ids, gaze_targets):
for f, id in zip(faces, ids):
# bbox_color = (255, 0, 0) if people[id].gender == 0 else (0, 0, 255)
bbox_color = (255, 0, 0)
draw_person_attributes(draw, people[id], f, f)
draw_bounding_box_pil(f, draw, bbox_color)
# draw_gaze_target_pil((int((g[0] - 1280) / 3 + 320), int(g[1] / 3)), draw, (0, 0, 255))
image_rgb = np.array(image_pil)
image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
cv2.imshow('video', image_bgr)
k = cv2.waitKey(1)
if k == 27: # Esc key to stop
break