7b89b8fc3f
Differential Revision: D90237984 fbshipit-source-id: 526fd760f303bf31be4f743bdcd77760496de0de
640 lines
25 KiB
Python
640 lines
25 KiB
Python
# fmt: off
|
|
# flake8: noqa
|
|
|
|
# pyre-unsafe
|
|
|
|
"""COCO Dataset."""
|
|
import copy
|
|
import itertools
|
|
import json
|
|
import os
|
|
from collections import defaultdict
|
|
|
|
import numpy as np
|
|
from scipy.optimize import linear_sum_assignment
|
|
|
|
from .. import _timing, utils
|
|
from ..config import get_default_dataset_config, init_config
|
|
from ..utils import TrackEvalException
|
|
from ._base_dataset import _BaseDataset
|
|
|
|
|
|
class COCO(_BaseDataset):
|
|
"""Tracking datasets in COCO format."""
|
|
|
|
def __init__(self, config=None):
|
|
"""Initialize dataset, checking that all required files are present."""
|
|
super().__init__()
|
|
# Fill non-given config values with defaults
|
|
self.config = init_config(config, get_default_dataset_config(), self.get_name())
|
|
self.gt_fol = self.config["GT_FOLDER"]
|
|
self.tracker_fol = self.config["TRACKERS_FOLDER"]
|
|
self.should_classes_combine = True
|
|
self.use_super_categories = False
|
|
self.use_mask = self.config["USE_MASK"]
|
|
|
|
self.tracker_sub_fol = self.config["TRACKER_SUB_FOLDER"]
|
|
self.output_fol = self.config["OUTPUT_FOLDER"]
|
|
if self.output_fol is None:
|
|
self.output_fol = self.tracker_fol
|
|
self.output_sub_fol = self.config["OUTPUT_SUB_FOLDER"]
|
|
|
|
if self.gt_fol.endswith(".json"):
|
|
self.gt_data = json.load(open(self.gt_fol, "r"))
|
|
else:
|
|
gt_dir_files = [
|
|
file for file in os.listdir(self.gt_fol) if file.endswith(".json")
|
|
]
|
|
if len(gt_dir_files) != 1:
|
|
raise TrackEvalException(
|
|
f"{self.gt_fol} does not contain exactly one json file."
|
|
)
|
|
|
|
with open(os.path.join(self.gt_fol, gt_dir_files[0])) as f:
|
|
self.gt_data = json.load(f)
|
|
|
|
# fill missing video ids
|
|
self._fill_video_ids_inplace(self.gt_data["annotations"])
|
|
|
|
# get sequences to eval and sequence information
|
|
self.seq_list = [
|
|
vid["name"].replace("/", "-") for vid in self.gt_data["videos"]
|
|
]
|
|
self.seq_name2seqid = {
|
|
vid["name"].replace("/", "-"): vid["id"] for vid in self.gt_data["videos"]
|
|
}
|
|
# compute mappings from videos to annotation data
|
|
self.video2gt_track, self.video2gt_image = self._compute_vid_mappings(
|
|
self.gt_data["annotations"]
|
|
)
|
|
# compute sequence lengths
|
|
self.seq_lengths = {vid["id"]: 0 for vid in self.gt_data["videos"]}
|
|
for img in self.gt_data["images"]:
|
|
self.seq_lengths[img["video_id"]] += 1
|
|
self.seq2images2timestep = self._compute_image_to_timestep_mappings()
|
|
self.seq2cls = {
|
|
vid["id"]: {
|
|
"pos_cat_ids": list(
|
|
{track["category_id"] for track in self.video2gt_track[vid["id"]]}
|
|
),
|
|
}
|
|
for vid in self.gt_data["videos"]
|
|
}
|
|
|
|
# Get classes to eval
|
|
considered_vid_ids = [self.seq_name2seqid[vid] for vid in self.seq_list]
|
|
seen_cats = set(
|
|
[
|
|
cat_id
|
|
for vid_id in considered_vid_ids
|
|
for cat_id in self.seq2cls[vid_id]["pos_cat_ids"]
|
|
]
|
|
)
|
|
# only classes with ground truth are evaluated in TAO
|
|
self.valid_classes = [
|
|
cls["name"] for cls in self.gt_data["categories"] if cls["id"] in seen_cats
|
|
]
|
|
cls_name2clsid_map = {
|
|
cls["name"]: cls["id"] for cls in self.gt_data["categories"]
|
|
}
|
|
|
|
if self.config["CLASSES_TO_EVAL"]:
|
|
self.class_list = [
|
|
cls.lower() if cls.lower() in self.valid_classes else None
|
|
for cls in self.config["CLASSES_TO_EVAL"]
|
|
]
|
|
if not all(self.class_list):
|
|
valid_cls = ", ".join(self.valid_classes)
|
|
raise TrackEvalException(
|
|
"Attempted to evaluate an invalid class. Only classes "
|
|
f"{valid_cls} are valid (classes present in ground truth"
|
|
" data)."
|
|
)
|
|
else:
|
|
self.class_list = [cls for cls in self.valid_classes]
|
|
self.cls_name2clsid = {
|
|
k: v for k, v in cls_name2clsid_map.items() if k in self.class_list
|
|
}
|
|
self.clsid2cls_name = {
|
|
v: k for k, v in cls_name2clsid_map.items() if k in self.class_list
|
|
}
|
|
# get trackers to eval
|
|
if self.config["TRACKERS_TO_EVAL"] is None:
|
|
self.tracker_list = os.listdir(self.tracker_fol)
|
|
else:
|
|
self.tracker_list = self.config["TRACKERS_TO_EVAL"]
|
|
|
|
if self.config["TRACKER_DISPLAY_NAMES"] is None:
|
|
self.tracker_to_disp = dict(zip(self.tracker_list, self.tracker_list))
|
|
elif (self.config["TRACKERS_TO_EVAL"] is not None) and (
|
|
len(self.config["TK_DISPLAY_NAMES"]) == len(self.tracker_list)
|
|
):
|
|
self.tracker_to_disp = dict(
|
|
zip(self.tracker_list, self.config["TK_DISPLAY_NAMES"])
|
|
)
|
|
else:
|
|
raise TrackEvalException(
|
|
"List of tracker files and tracker display names do not match."
|
|
)
|
|
|
|
self.tracker_data = {tracker: dict() for tracker in self.tracker_list}
|
|
|
|
for tracker in self.tracker_list:
|
|
if self.tracker_sub_fol.endswith(".json"):
|
|
with open(os.path.join(self.tracker_sub_fol)) as f:
|
|
curr_data = json.load(f)
|
|
else:
|
|
tr_dir = os.path.join(self.tracker_fol, tracker, self.tracker_sub_fol)
|
|
tr_dir_files = [
|
|
file for file in os.listdir(tr_dir) if file.endswith(".json")
|
|
]
|
|
if len(tr_dir_files) != 1:
|
|
raise TrackEvalException(
|
|
f"{tr_dir} does not contain exactly one json file."
|
|
)
|
|
with open(os.path.join(tr_dir, tr_dir_files[0])) as f:
|
|
curr_data = json.load(f)
|
|
|
|
# limit detections if MAX_DETECTIONS > 0
|
|
if self.config["MAX_DETECTIONS"]:
|
|
curr_data = self._limit_dets_per_image(curr_data)
|
|
|
|
# fill missing video ids
|
|
self._fill_video_ids_inplace(curr_data)
|
|
|
|
# make track ids unique over whole evaluation set
|
|
self._make_tk_ids_unique(curr_data)
|
|
|
|
# get tracker sequence information
|
|
curr_vids2tracks, curr_vids2images = self._compute_vid_mappings(curr_data)
|
|
self.tracker_data[tracker]["vids_to_tracks"] = curr_vids2tracks
|
|
self.tracker_data[tracker]["vids_to_images"] = curr_vids2images
|
|
|
|
def get_display_name(self, tracker):
|
|
return self.tracker_to_disp[tracker]
|
|
|
|
def _load_raw_file(self, tracker, seq, is_gt):
|
|
"""Load a file (gt or tracker) in the TAO format
|
|
|
|
If is_gt, this returns a dict which contains the fields:
|
|
[gt_ids, gt_classes]:
|
|
list (for each timestep) of 1D NDArrays (for each det).
|
|
[gt_dets]: list (for each timestep) of lists of detections.
|
|
|
|
if not is_gt, this returns a dict which contains the fields:
|
|
[tk_ids, tk_classes]:
|
|
list (for each timestep) of 1D NDArrays (for each det).
|
|
[tk_dets]: list (for each timestep) of lists of detections.
|
|
"""
|
|
seq_id = self.seq_name2seqid[seq]
|
|
# file location
|
|
if is_gt:
|
|
imgs = self.video2gt_image[seq_id]
|
|
else:
|
|
imgs = self.tracker_data[tracker]["vids_to_images"][seq_id]
|
|
|
|
# convert data to required format
|
|
num_timesteps = self.seq_lengths[seq_id]
|
|
img_to_timestep = self.seq2images2timestep[seq_id]
|
|
data_keys = ["ids", "classes", "dets"]
|
|
# if not is_gt:
|
|
# data_keys += ["tk_confidences"]
|
|
raw_data = {key: [None] * num_timesteps for key in data_keys}
|
|
for img in imgs:
|
|
# some tracker data contains images without any ground truth info,
|
|
# these are ignored
|
|
if img["id"] not in img_to_timestep:
|
|
continue
|
|
t = img_to_timestep[img["id"]]
|
|
anns = img["annotations"]
|
|
tk_str = utils.get_track_id_str(anns[0])
|
|
|
|
if self.use_mask:
|
|
# When using mask, extract segmentation data
|
|
raw_data["dets"][t] = [ann.get("segmentation") for ann in anns]
|
|
else:
|
|
# When using bbox, extract bbox data
|
|
raw_data["dets"][t] = np.atleast_2d([ann["bbox"] for ann in anns]).astype(
|
|
float
|
|
)
|
|
raw_data["ids"][t] = np.atleast_1d([ann[tk_str] for ann in anns]).astype(
|
|
int
|
|
)
|
|
raw_data["classes"][t] = np.atleast_1d(
|
|
[ann["category_id"] for ann in anns]
|
|
).astype(int)
|
|
# if not is_gt:
|
|
# raw_data["tk_confidences"][t] = np.atleast_1d(
|
|
# [ann["score"] for ann in anns]
|
|
# ).astype(float)
|
|
|
|
for t, d in enumerate(raw_data["dets"]):
|
|
if d is None:
|
|
raw_data["dets"][t] = np.empty((0, 4)).astype(float)
|
|
raw_data["ids"][t] = np.empty(0).astype(int)
|
|
raw_data["classes"][t] = np.empty(0).astype(int)
|
|
# if not is_gt:
|
|
# raw_data["tk_confidences"][t] = np.empty(0)
|
|
|
|
if is_gt:
|
|
key_map = {"ids": "gt_ids", "classes": "gt_classes", "dets": "gt_dets"}
|
|
else:
|
|
key_map = {"ids": "tk_ids", "classes": "tk_classes", "dets": "tk_dets"}
|
|
for k, v in key_map.items():
|
|
raw_data[v] = raw_data.pop(k)
|
|
|
|
raw_data["num_timesteps"] = num_timesteps
|
|
raw_data["seq"] = seq
|
|
return raw_data
|
|
|
|
def get_preprocessed_seq_data_thr(self, raw_data, cls, assignment=None):
|
|
"""Preprocess data for a single sequence for a single class.
|
|
|
|
Inputs:
|
|
raw_data: dict containing the data for the sequence already
|
|
read in by get_raw_seq_data().
|
|
cls: class to be evaluated.
|
|
Outputs:
|
|
gt_ids:
|
|
list (for each timestep) of ids of GT tracks
|
|
tk_ids:
|
|
list (for each timestep) of ids of predicted tracks (all for TP
|
|
matching (Det + AssocA))
|
|
tk_overlap_ids:
|
|
list (for each timestep) of ids of predicted tracks that overlap
|
|
with GTs
|
|
tk_dets:
|
|
list (for each timestep) of lists of detections that
|
|
corresponding to the tk_ids
|
|
tk_classes:
|
|
list (for each timestep) of lists of classes that corresponding
|
|
to the tk_ids
|
|
tk_confidences:
|
|
list (for each timestep) of lists of classes that corresponding
|
|
to the tk_ids
|
|
sim_scores:
|
|
similarity score between gt_ids and tk_ids.
|
|
"""
|
|
if cls != "all":
|
|
cls_id = self.cls_name2clsid[cls]
|
|
|
|
data_keys = [
|
|
"gt_ids",
|
|
"tk_ids",
|
|
"gt_id_map",
|
|
"tk_id_map",
|
|
"gt_dets",
|
|
"gt_classes",
|
|
"gt_class_name",
|
|
"tk_overlap_classes",
|
|
"tk_overlap_ids",
|
|
"tk_class_eval_tk_ids",
|
|
"tk_dets",
|
|
"tk_classes",
|
|
# "tk_confidences",
|
|
"tk_exh_ids",
|
|
"sim_scores",
|
|
]
|
|
data = {key: [None] * raw_data["num_timesteps"] for key in data_keys}
|
|
unique_gt_ids = []
|
|
unique_tk_ids = []
|
|
num_gt_dets = 0
|
|
num_tk_cls_dets = 0
|
|
num_tk_overlap_dets = 0
|
|
overlap_ious_thr = 0.5
|
|
loc_and_asso_tk_ids = []
|
|
exh_class_tk_ids = []
|
|
|
|
for t in range(raw_data["num_timesteps"]):
|
|
# only extract relevant dets for this class for preproc and eval
|
|
if cls == "all":
|
|
gt_class_mask = np.ones_like(raw_data["gt_classes"][t]).astype(bool)
|
|
else:
|
|
gt_class_mask = np.atleast_1d(
|
|
raw_data["gt_classes"][t] == cls_id
|
|
).astype(bool)
|
|
|
|
# select GT that is not in the evaluating classes
|
|
if assignment is not None and assignment:
|
|
all_gt_ids = list(assignment[t].keys())
|
|
gt_ids_in = raw_data["gt_ids"][t][gt_class_mask]
|
|
gt_ids_out = set(all_gt_ids) - set(gt_ids_in)
|
|
tk_ids_out = set([assignment[t][key] for key in list(gt_ids_out)])
|
|
|
|
# compute overlapped tracks and add their ids to overlap_tk_ids
|
|
sim_scores = raw_data["similarity_scores"]
|
|
overlap_ids_masks = (sim_scores[t][gt_class_mask] >= overlap_ious_thr).any(
|
|
axis=0
|
|
)
|
|
overlap_tk_ids_t = raw_data["tk_ids"][t][overlap_ids_masks]
|
|
if assignment is not None and assignment:
|
|
data["tk_overlap_ids"][t] = list(set(overlap_tk_ids_t) - tk_ids_out)
|
|
else:
|
|
data["tk_overlap_ids"][t] = list(set(overlap_tk_ids_t))
|
|
|
|
loc_and_asso_tk_ids += data["tk_overlap_ids"][t]
|
|
|
|
data["tk_exh_ids"][t] = []
|
|
if cls == "all":
|
|
continue
|
|
|
|
# add the track ids of exclusive annotated class to exh_class_tk_ids
|
|
tk_exh_mask = np.atleast_1d(raw_data["tk_classes"][t] == cls_id)
|
|
tk_exh_mask = tk_exh_mask.astype(bool)
|
|
exh_class_tk_ids_t = raw_data["tk_ids"][t][tk_exh_mask]
|
|
exh_class_tk_ids.append(exh_class_tk_ids_t)
|
|
data["tk_exh_ids"][t] = exh_class_tk_ids_t
|
|
|
|
# remove tk_ids that has been assigned to GT belongs to other classes.
|
|
loc_and_asso_tk_ids = list(set(loc_and_asso_tk_ids))
|
|
|
|
# remove all unwanted unmatched tracker detections
|
|
for t in range(raw_data["num_timesteps"]):
|
|
# add gt to the data
|
|
if cls == "all":
|
|
gt_class_mask = np.ones_like(raw_data["gt_classes"][t]).astype(bool)
|
|
else:
|
|
gt_class_mask = np.atleast_1d(
|
|
raw_data["gt_classes"][t] == cls_id
|
|
).astype(bool)
|
|
data["gt_classes"][t] = cls_id
|
|
data["gt_class_name"][t] = cls
|
|
|
|
gt_ids = raw_data["gt_ids"][t][gt_class_mask]
|
|
if self.use_mask:
|
|
gt_dets = [raw_data['gt_dets'][t][ind] for ind in range(len(gt_class_mask)) if gt_class_mask[ind]]
|
|
else:
|
|
gt_dets = raw_data["gt_dets"][t][gt_class_mask]
|
|
data["gt_ids"][t] = gt_ids
|
|
data["gt_dets"][t] = gt_dets
|
|
|
|
# filter pred and only keep those that highly overlap with GTs
|
|
tk_mask = np.isin(
|
|
raw_data["tk_ids"][t], np.array(loc_and_asso_tk_ids), assume_unique=True
|
|
)
|
|
tk_overlap_mask = np.isin(
|
|
raw_data["tk_ids"][t],
|
|
np.array(data["tk_overlap_ids"][t]),
|
|
assume_unique=True,
|
|
)
|
|
|
|
tk_ids = raw_data["tk_ids"][t][tk_mask]
|
|
if self.use_mask:
|
|
tk_dets = [raw_data['tk_dets'][t][ind] for ind in range(len(tk_mask)) if
|
|
tk_mask[ind]]
|
|
else:
|
|
tk_dets = raw_data["tk_dets"][t][tk_mask]
|
|
|
|
tracker_classes = raw_data["tk_classes"][t][tk_mask]
|
|
|
|
# add overlap classes for computing the FP for Cls term
|
|
tracker_overlap_classes = raw_data["tk_classes"][t][tk_overlap_mask]
|
|
# tracker_confidences = raw_data["tk_confidences"][t][tk_mask]
|
|
sim_scores_masked = sim_scores[t][gt_class_mask, :][:, tk_mask]
|
|
|
|
# add filtered prediction to the data
|
|
data["tk_classes"][t] = tracker_classes
|
|
data["tk_overlap_classes"][t] = tracker_overlap_classes
|
|
data["tk_ids"][t] = tk_ids
|
|
data["tk_dets"][t] = tk_dets
|
|
# data["tk_confidences"][t] = tracker_confidences
|
|
data["sim_scores"][t] = sim_scores_masked
|
|
data["tk_class_eval_tk_ids"][t] = set(
|
|
list(data["tk_overlap_ids"][t]) + list(data["tk_exh_ids"][t])
|
|
)
|
|
|
|
# count total number of detections
|
|
unique_gt_ids += list(np.unique(data["gt_ids"][t]))
|
|
# the unique track ids are for association.
|
|
unique_tk_ids += list(np.unique(data["tk_ids"][t]))
|
|
|
|
num_tk_overlap_dets += len(data["tk_overlap_ids"][t])
|
|
num_tk_cls_dets += len(data["tk_class_eval_tk_ids"][t])
|
|
num_gt_dets += len(data["gt_ids"][t])
|
|
|
|
# re-label IDs such that there are no empty IDs
|
|
if len(unique_gt_ids) > 0:
|
|
unique_gt_ids = np.unique(unique_gt_ids)
|
|
gt_id_map = np.nan * np.ones((np.max(unique_gt_ids) + 1))
|
|
gt_id_map[unique_gt_ids] = np.arange(len(unique_gt_ids))
|
|
data["gt_id_map"] = {}
|
|
for gt_id in unique_gt_ids:
|
|
new_gt_id = gt_id_map[gt_id].astype(int)
|
|
data["gt_id_map"][new_gt_id] = gt_id
|
|
|
|
for t in range(raw_data["num_timesteps"]):
|
|
if len(data["gt_ids"][t]) > 0:
|
|
data["gt_ids"][t] = gt_id_map[data["gt_ids"][t]].astype(int)
|
|
|
|
if len(unique_tk_ids) > 0:
|
|
unique_tk_ids = np.unique(unique_tk_ids)
|
|
tk_id_map = np.nan * np.ones((np.max(unique_tk_ids) + 1))
|
|
tk_id_map[unique_tk_ids] = np.arange(len(unique_tk_ids))
|
|
|
|
data["tk_id_map"] = {}
|
|
for track_id in unique_tk_ids:
|
|
new_track_id = tk_id_map[track_id].astype(int)
|
|
data["tk_id_map"][new_track_id] = track_id
|
|
|
|
for t in range(raw_data["num_timesteps"]):
|
|
if len(data["tk_ids"][t]) > 0:
|
|
data["tk_ids"][t] = tk_id_map[data["tk_ids"][t]].astype(int)
|
|
if len(data["tk_overlap_ids"][t]) > 0:
|
|
data["tk_overlap_ids"][t] = tk_id_map[
|
|
data["tk_overlap_ids"][t]
|
|
].astype(int)
|
|
|
|
# record overview statistics.
|
|
data["num_tk_cls_dets"] = num_tk_cls_dets
|
|
data["num_tk_overlap_dets"] = num_tk_overlap_dets
|
|
data["num_gt_dets"] = num_gt_dets
|
|
data["num_tk_ids"] = len(unique_tk_ids)
|
|
data["num_gt_ids"] = len(unique_gt_ids)
|
|
data["num_timesteps"] = raw_data["num_timesteps"]
|
|
data["seq"] = raw_data["seq"]
|
|
|
|
self._check_unique_ids(data)
|
|
|
|
return data
|
|
|
|
@_timing.time
|
|
def get_preprocessed_seq_data(
|
|
self, raw_data, cls, assignment=None, thresholds=[50, 75]
|
|
):
|
|
"""Preprocess data for a single sequence for a single class."""
|
|
data = {}
|
|
if thresholds is None:
|
|
thresholds = [50, 75]
|
|
elif isinstance(thresholds, int):
|
|
thresholds = [thresholds]
|
|
|
|
for thr in thresholds:
|
|
assignment_thr = None
|
|
if assignment is not None:
|
|
assignment_thr = assignment[thr]
|
|
data[thr] = self.get_preprocessed_seq_data_thr(
|
|
raw_data, cls, assignment_thr
|
|
)
|
|
|
|
return data
|
|
|
|
def _calculate_similarities(self, gt_dets_t, tk_dets_t):
|
|
"""Compute similarity scores."""
|
|
if self.use_mask:
|
|
similarity_scores = self._calculate_mask_ious(gt_dets_t, tk_dets_t, is_encoded=True, do_ioa=False)
|
|
else:
|
|
similarity_scores = self._calculate_box_ious(gt_dets_t, tk_dets_t)
|
|
return similarity_scores
|
|
|
|
def _compute_vid_mappings(self, annotations):
|
|
"""Computes mappings from videos to corresponding tracks and images."""
|
|
vids_to_tracks = {}
|
|
vids_to_imgs = {}
|
|
vid_ids = [vid["id"] for vid in self.gt_data["videos"]]
|
|
|
|
# compute an mapping from image IDs to images
|
|
images = {}
|
|
for image in self.gt_data["images"]:
|
|
images[image["id"]] = image
|
|
|
|
tk_str = utils.get_track_id_str(annotations[0])
|
|
for ann in annotations:
|
|
ann["area"] = ann["bbox"][2] * ann["bbox"][3]
|
|
|
|
vid = ann["video_id"]
|
|
if ann["video_id"] not in vids_to_tracks.keys():
|
|
vids_to_tracks[ann["video_id"]] = list()
|
|
if ann["video_id"] not in vids_to_imgs.keys():
|
|
vids_to_imgs[ann["video_id"]] = list()
|
|
|
|
# fill in vids_to_tracks
|
|
tid = ann[tk_str]
|
|
exist_tids = [track["id"] for track in vids_to_tracks[vid]]
|
|
try:
|
|
index1 = exist_tids.index(tid)
|
|
except ValueError:
|
|
index1 = -1
|
|
if tid not in exist_tids:
|
|
curr_track = {
|
|
"id": tid,
|
|
"category_id": ann["category_id"],
|
|
"video_id": vid,
|
|
"annotations": [ann],
|
|
}
|
|
vids_to_tracks[vid].append(curr_track)
|
|
else:
|
|
vids_to_tracks[vid][index1]["annotations"].append(ann)
|
|
|
|
# fill in vids_to_imgs
|
|
img_id = ann["image_id"]
|
|
exist_img_ids = [img["id"] for img in vids_to_imgs[vid]]
|
|
try:
|
|
index2 = exist_img_ids.index(img_id)
|
|
except ValueError:
|
|
index2 = -1
|
|
if index2 == -1:
|
|
curr_img = {"id": img_id, "annotations": [ann]}
|
|
vids_to_imgs[vid].append(curr_img)
|
|
else:
|
|
vids_to_imgs[vid][index2]["annotations"].append(ann)
|
|
|
|
# sort annotations by frame index and compute track area
|
|
for vid, tracks in vids_to_tracks.items():
|
|
for track in tracks:
|
|
track["annotations"] = sorted(
|
|
track["annotations"],
|
|
key=lambda x: images[x["image_id"]]["frame_id"],
|
|
)
|
|
# compute average area
|
|
track["area"] = sum(x["area"] for x in track["annotations"]) / len(
|
|
track["annotations"]
|
|
)
|
|
|
|
# ensure all videos are present
|
|
for vid_id in vid_ids:
|
|
if vid_id not in vids_to_tracks.keys():
|
|
vids_to_tracks[vid_id] = []
|
|
if vid_id not in vids_to_imgs.keys():
|
|
vids_to_imgs[vid_id] = []
|
|
|
|
return vids_to_tracks, vids_to_imgs
|
|
|
|
def _compute_image_to_timestep_mappings(self):
|
|
"""Computes a mapping from images to timestep in sequence."""
|
|
images = {}
|
|
for image in self.gt_data["images"]:
|
|
images[image["id"]] = image
|
|
|
|
seq_to_imgs_to_timestep = {vid["id"]: dict() for vid in self.gt_data["videos"]}
|
|
for vid in seq_to_imgs_to_timestep:
|
|
curr_imgs = [img["id"] for img in self.video2gt_image[vid]]
|
|
curr_imgs = sorted(curr_imgs, key=lambda x: images[x]["frame_id"])
|
|
seq_to_imgs_to_timestep[vid] = {
|
|
curr_imgs[i]: i for i in range(len(curr_imgs))
|
|
}
|
|
|
|
return seq_to_imgs_to_timestep
|
|
|
|
def _limit_dets_per_image(self, annotations):
|
|
"""Limits the number of detections for each image.
|
|
|
|
Adapted from https://github.com/TAO-Dataset/.
|
|
"""
|
|
max_dets = self.config["MAX_DETECTIONS"]
|
|
img_ann = defaultdict(list)
|
|
for ann in annotations:
|
|
img_ann[ann["image_id"]].append(ann)
|
|
|
|
for img_id, _anns in img_ann.items():
|
|
if len(_anns) <= max_dets:
|
|
continue
|
|
_anns = sorted(_anns, key=lambda x: x["score"], reverse=True)
|
|
img_ann[img_id] = _anns[:max_dets]
|
|
|
|
return [ann for anns in img_ann.values() for ann in anns]
|
|
|
|
def _fill_video_ids_inplace(self, annotations):
|
|
"""Fills in missing video IDs inplace.
|
|
|
|
Adapted from https://github.com/TAO-Dataset/.
|
|
"""
|
|
missing_video_id = [x for x in annotations if "video_id" not in x]
|
|
if missing_video_id:
|
|
image_id_to_video_id = {
|
|
x["id"]: x["video_id"] for x in self.gt_data["images"]
|
|
}
|
|
for x in missing_video_id:
|
|
x["video_id"] = image_id_to_video_id[x["image_id"]]
|
|
|
|
@staticmethod
|
|
def _make_tk_ids_unique(annotations):
|
|
"""Makes track IDs unqiue over the whole annotation set.
|
|
|
|
Adapted from https://github.com/TAO-Dataset/.
|
|
"""
|
|
track_id_videos = {}
|
|
track_ids_to_update = set()
|
|
max_track_id = 0
|
|
|
|
tk_str = utils.get_track_id_str(annotations[0])
|
|
for ann in annotations:
|
|
t = int(ann[tk_str])
|
|
if t not in track_id_videos:
|
|
track_id_videos[t] = ann["video_id"]
|
|
|
|
if ann["video_id"] != track_id_videos[t]:
|
|
# track id is assigned to multiple videos
|
|
track_ids_to_update.add(t)
|
|
max_track_id = max(max_track_id, t)
|
|
|
|
if track_ids_to_update:
|
|
print("true")
|
|
next_id = itertools.count(max_track_id + 1)
|
|
new_tk_ids = defaultdict(lambda: next(next_id))
|
|
for ann in annotations:
|
|
t = ann[tk_str]
|
|
v = ann["video_id"]
|
|
if t in track_ids_to_update:
|
|
ann[tk_str] = new_tk_ids[t, v]
|
|
return len(track_ids_to_update)
|