def track(self, boxes, labels, probs, box_feats, frame_num):
"""Given new object detection output, update the tracking."""
obj_id_to_class_ = obj_id_to_class
if self.is_coco_class:
obj_id_to_class_ = coco_obj_id_to_class
if self.partial_classes:
obj_id_to_class_ = self.partial_obj_id2class
detections = create_obj_infos(
frame_num,
boxes,
probs,
labels,
box_feats,
obj_id_to_class_, [self.track_obj],
self.min_confidence,
self.min_detection_height,
1.0,
is_coco_model=self.is_coco_class,
coco_to_actev_mapping=coco_obj_to_actev_obj)
# Run non-maxima suppression.
boxes = np.array([d.tlwh for d in detections])
scores = np.array([d.confidence for d in detections])
indices = preprocessing.non_max_suppression(boxes,
self.nms_max_overlap,
scores)
detections = [detections[i] for i in indices]
# update tracker
self.tracker.predict()
self.tracker.update(detections)
# store result
for track in self.tracker.tracks:
if not track.is_confirmed() or track.time_since_update > 1:
if not track.is_confirmed and track.time_since_update == 0:
bbox = track.to_tlwh()
dp = [
frame_num, track.track_id, bbox[0], bbox[1], bbox[2],
bbox[3]
]
if track.track_id not in self.tmp_tracking_dict:
self.tmp_tracking_dict[track.track_id] = [dp]
else:
self.tmp_tracking_dict[track.track_id].append(dp)
continue
bbox = track.to_tlwh()
if track.track_id in self.tmp_tracking_dict:
pred_list = self.tmp_tracking_dict[track.track_id]
for dp in pred_list:
self.tracking_results.append(dp)
self.tmp_tracking_dict[track.track_id].pop(
track.track_id, None)
self.tracking_results.append([
frame_num, track.track_id, bbox[0], bbox[1], bbox[2], bbox[3]
])
def run_detect_and_track(args,
frame_stack,
sess,
model,
targetid2class,
tracking_objs,
tracker_dict,
tracking_results_dict,
tmp_tracking_results_dict,
obj_out_dir=None,
valid_frame_num=None):
# ignore the padded images
if valid_frame_num is None:
valid_frame_num = len(frame_stack)
resized_images, scales, frame_idxs = zip(*frame_stack)
feed_dict = model.get_feed_dict_forward_multi(resized_images)
sess_input = [
model.final_boxes, model.final_labels, model.final_probs,
model.final_valid_indices, model.fpn_box_feat
]
# [B, num, 4], [B, num], [B, num], [B], [M, 256, 7, 7]
batch_boxes, batch_labels, batch_probs, valid_indices, batch_box_feats = \
sess.run(sess_input, feed_dict=feed_dict)
assert np.sum(valid_indices) == batch_box_feats.shape[0], "duh"
for b in range(valid_frame_num):
cur_frame = frame_idxs[b]
# [k, 4]
final_boxes = batch_boxes[b][:valid_indices[b]]
# [k]
final_labels = batch_labels[b][:valid_indices[b]]
# [k]
final_probs = batch_probs[b][:valid_indices[b]]
# [k, 256, 7, 7]
previous_box_num = sum(valid_indices[:b])
box_feats = batch_box_feats[previous_box_num:previous_box_num +
valid_indices[b]]
if args.get_tracking:
assert len(box_feats) == len(final_boxes)
for tracking_obj in tracking_objs:
target_tracking_obs = [tracking_obj]
# will consider scale here
scale = scales[b]
detections = create_obj_infos(
cur_frame,
final_boxes,
final_probs,
final_labels,
box_feats,
targetid2class,
target_tracking_obs,
args.min_confidence,
args.min_detection_height,
scale,
is_coco_model=args.is_coco_model,
coco_to_actev_mapping=coco_obj_to_actev_obj)
# Run non-maxima suppression.
boxes = np.array([d.tlwh for d in detections])
scores = np.array([d.confidence for d in detections])
indices = preprocessing.non_max_suppression(
boxes, args.nms_max_overlap, scores)
detections = [detections[i] for i in indices]
# tracking
tracker_dict[tracking_obj].predict()
tracker_dict[tracking_obj].update(detections)
# Store results
for track in tracker_dict[tracking_obj].tracks:
if not track.is_confirmed() or track.time_since_update > 1:
if (not track.is_confirmed()
) and track.time_since_update == 0:
bbox = track.to_tlwh()
if track.track_id not in \
tmp_tracking_results_dict[tracking_obj]:
tmp_tracking_results_dict[tracking_obj][track.track_id] = \
[[cur_frame, track.track_id, bbox[0], bbox[1],
bbox[2], bbox[3]]]
else:
tmp_tracking_results_dict[tracking_obj][
track.track_id].append([
cur_frame, track.track_id, bbox[0],
bbox[1], bbox[2], bbox[3]
])
continue
bbox = track.to_tlwh()
if track.track_id in tmp_tracking_results_dict[
tracking_obj]:
pred_list = tmp_tracking_results_dict[tracking_obj][
track.track_id]
for pred_data in pred_list:
tracking_results_dict[tracking_obj].append(
pred_data)
tmp_tracking_results_dict[tracking_obj].pop(
track.track_id, None)
tracking_results_dict[tracking_obj].append([
cur_frame, track.track_id, bbox[0], bbox[1], bbox[2],
bbox[3]
])
if obj_out_dir is None: # not saving the boxes
continue
# ---------------- get the json outputs for object detection
# scale back the box to original image size
final_boxes = final_boxes / scales[b]
# save as json
pred = []
for j, (box, prob,
label) in enumerate(zip(final_boxes, final_probs,
final_labels)):
box[2] -= box[0]
box[3] -= box[1] # produce x,y,w,h output
cat_id = int(label)
cat_name = targetid2class[cat_id]
res = {
"category_id": int(cat_id),
"cat_name": cat_name, # [0-80]
"score": float(round(prob, 7)),
#"bbox": list(map(lambda x: float(round(x, 2)), box)),
"bbox": [float(round(x, 2)) for x in box],
"segmentation": None,
}
pred.append(res)
predfile = os.path.join(video_out_path, "%d.json" % (cur_frame))
with open(predfile, "w") as f:
json.dump(pred, f)
model.final_probs, model.fpn_box_feat
]
final_boxes, final_labels, final_probs, box_feats = sess.run(
sess_input, feed_dict=feed_dict)
assert len(box_feats) == len(final_boxes)
for tracking_obj in tracking_objs:
target_tracking_obs = [tracking_obj]
detections = create_obj_infos(
cur_frame,
final_boxes,
final_probs,
final_labels,
box_feats,
targetid2class,
target_tracking_obs,
args.min_confidence,
args.min_detection_height,
scale,
is_coco_model=args.is_coco_model,
coco_to_actev_mapping=coco_obj_to_actev_obj)
# Run non-maxima suppression.
boxes = np.array([d.tlwh for d in detections])
scores = np.array([d.confidence for d in detections])
indices = preprocessing.non_max_suppression(
boxes, args.nms_max_overlap, scores)
detections = [detections[i] for i in indices]
# tracking
tracker_dict[tracking_obj].predict()
"%d.npy" % (cur_frame))
np.save(featfile, box_feats)
elif args.get_tracking:
sess_input = [
model.final_boxes, model.final_labels,
model.final_probs, model.fpn_box_feat
]
final_boxes, final_labels, final_probs, box_feats = sess.run(
sess_input, feed_dict=feed_dict)
assert len(box_feats) == len(final_boxes)
for tracking_obj in tracking_objs:
target_tracking_obs = [tracking_obj]
detections = create_obj_infos(
cur_frame, final_boxes, final_probs, final_labels,
box_feats, targetid2class, target_tracking_obs,
args.min_confidence, args.min_detection_height,
scale)
# Run non-maxima suppression.
boxes = np.array([d.tlwh for d in detections])
scores = np.array([d.confidence for d in detections])
indices = preprocessing.non_max_suppression(
boxes, args.nms_max_overlap, scores)
detections = [detections[i] for i in indices]
# tracking
tracker_dict[tracking_obj].predict()
tracker_dict[tracking_obj].update(detections)
# Store results
for track in tracker_dict[tracking_obj].tracks: