|
- # Copyright 2022 Huawei Technologies Co., Ltd
- #
- # Licensed under the Apache License, Version 2.0 (the "License");
- # you may not use this file except in compliance with the License.
- # You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
- # ============================================================================
- """MOT utils."""
- import os
-
- import numpy as np
-
-
- def read_results(filename, data_type: str, is_gt=False, is_ignore=False):
- """
- Read results.
- """
- if data_type in ('mot', 'lab'):
- read_fun = read_mot_results
- else:
- raise ValueError('Unknown data type: {data_type}')
-
- return read_fun(filename, is_gt, is_ignore)
-
-
- def read_mot_results(filename, is_gt, is_ignore):
- """
- Read MOT results.
- """
- valid_labels = {1}
- ignore_labels = {2, 7, 8, 12}
- results_dict = {}
- if os.path.isfile(filename):
- with open(filename, 'r') as f:
- for line in f.readlines():
- linelist = line.split(',')
- if len(linelist) < 7:
- continue
- fid = int(linelist[0])
- if fid < 1:
- continue
- results_dict.setdefault(fid, [])
-
- if is_gt:
- if 'MOT16-' in filename or 'MOT17-' in filename:
- label = int(float(linelist[7]))
- mark = int(float(linelist[6]))
- if mark == 0 or label not in valid_labels:
- continue
- score = 1
- elif is_ignore:
- if 'MOT16-' in filename or 'MOT17-' in filename:
- label = int(float(linelist[7]))
- vis_ratio = float(linelist[8])
- if label not in ignore_labels and vis_ratio >= 0:
- continue
- else:
- continue
- score = 1
- else:
- score = float(linelist[6])
-
- tlwh = tuple(map(float, linelist[2:6]))
- target_id = int(linelist[1])
-
- results_dict[fid].append((tlwh, target_id, score))
-
- return results_dict
-
-
- def unzip_objs(objs):
- """
- Unzip objects.
- """
- if objs:
- tlwhs, ids, scores = zip(*objs)
- else:
- tlwhs, ids, scores = [], [], []
- tlwhs = np.asarray(tlwhs, dtype=float).reshape(-1, 4)
-
- return tlwhs, ids, scores
|