|
|
|
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
|
|
|
#
|
|
|
|
# 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.
|
|
|
|
|
|
|
|
import os
|
|
|
|
import sys
|
|
|
|
import cv2
|
|
|
|
import glob
|
|
|
|
import numpy as np
|
|
|
|
from collections import OrderedDict, defaultdict
|
|
|
|
try:
|
|
|
|
from collections.abc import Sequence
|
|
|
|
except Exception:
|
|
|
|
from collections import Sequence
|
|
|
|
from .dataset import DetDataset, _make_dataset, _is_valid_file
|
|
|
|
from paddlers.models.ppdet.core.workspace import register, serializable
|
|
|
|
from paddlers.models.ppdet.utils.logger import setup_logger
|
|
|
|
logger = setup_logger(__name__)
|
|
|
|
|
|
|
|
|
|
|
|
@register
|
|
|
|
@serializable
|
|
|
|
class MOTDataSet(DetDataset):
|
|
|
|
"""
|
|
|
|
Load dataset with MOT format, only support single class MOT.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
dataset_dir (str): root directory for dataset.
|
|
|
|
image_lists (str|list): mot data image lists, muiti-source mot dataset.
|
|
|
|
data_fields (list): key name of data dictionary, at least have 'image'.
|
|
|
|
sample_num (int): number of samples to load, -1 means all.
|
|
|
|
|
|
|
|
Notes:
|
|
|
|
MOT datasets root directory following this:
|
|
|
|
dataset/mot
|
|
|
|
|——————image_lists
|
|
|
|
| |——————caltech.train
|
|
|
|
| |——————caltech.val
|
|
|
|
| |——————mot16.train
|
|
|
|
| |——————mot17.train
|
|
|
|
| ......
|
|
|
|
|——————Caltech
|
|
|
|
|——————MOT17
|
|
|
|
|——————......
|
|
|
|
|
|
|
|
All the MOT datasets have the following structure:
|
|
|
|
Caltech
|
|
|
|
|——————images
|
|
|
|
| └——————00001.jpg
|
|
|
|
| |—————— ...
|
|
|
|
| └——————0000N.jpg
|
|
|
|
└——————labels_with_ids
|
|
|
|
└——————00001.txt
|
|
|
|
|—————— ...
|
|
|
|
└——————0000N.txt
|
|
|
|
or
|
|
|
|
|
|
|
|
MOT17
|
|
|
|
|——————images
|
|
|
|
| └——————train
|
|
|
|
| └——————test
|
|
|
|
└——————labels_with_ids
|
|
|
|
└——————train
|
|
|
|
"""
|
|
|
|
|
|
|
|
def __init__(self,
|
|
|
|
dataset_dir=None,
|
|
|
|
image_lists=[],
|
|
|
|
data_fields=['image'],
|
|
|
|
sample_num=-1):
|
|
|
|
super(MOTDataSet, self).__init__(
|
|
|
|
dataset_dir=dataset_dir,
|
|
|
|
data_fields=data_fields,
|
|
|
|
sample_num=sample_num)
|
|
|
|
self.dataset_dir = dataset_dir
|
|
|
|
self.image_lists = image_lists
|
|
|
|
if isinstance(self.image_lists, str):
|
|
|
|
self.image_lists = [self.image_lists]
|
|
|
|
self.roidbs = None
|
|
|
|
self.cname2cid = None
|
|
|
|
|
|
|
|
def get_anno(self):
|
|
|
|
if self.image_lists == []:
|
|
|
|
return
|
|
|
|
# only used to get categories and metric
|
|
|
|
# only check first data, but the label_list of all data should be same.
|
|
|
|
first_mot_data = self.image_lists[0].split('.')[0]
|
|
|
|
anno_file = os.path.join(self.dataset_dir, first_mot_data,
|
|
|
|
'label_list.txt')
|
|
|
|
return anno_file
|
|
|
|
|
|
|
|
def parse_dataset(self):
|
|
|
|
self.img_files = OrderedDict()
|
|
|
|
self.img_start_index = OrderedDict()
|
|
|
|
self.label_files = OrderedDict()
|
|
|
|
self.tid_num = OrderedDict()
|
|
|
|
self.tid_start_index = OrderedDict()
|
|
|
|
|
|
|
|
img_index = 0
|
|
|
|
for data_name in self.image_lists:
|
|
|
|
# check every data image list
|
|
|
|
image_lists_dir = os.path.join(self.dataset_dir, 'image_lists')
|
|
|
|
assert os.path.isdir(image_lists_dir), \
|
|
|
|
"The {} is not a directory.".format(image_lists_dir)
|
|
|
|
|
|
|
|
list_path = os.path.join(image_lists_dir, data_name)
|
|
|
|
assert os.path.exists(list_path), \
|
|
|
|
"The list path {} does not exist.".format(list_path)
|
|
|
|
|
|
|
|
# record img_files, filter out empty ones
|
|
|
|
with open(list_path, 'r') as file:
|
|
|
|
self.img_files[data_name] = file.readlines()
|
|
|
|
self.img_files[data_name] = [
|
|
|
|
os.path.join(self.dataset_dir, x.strip())
|
|
|
|
for x in self.img_files[data_name]
|
|
|
|
]
|
|
|
|
self.img_files[data_name] = list(
|
|
|
|
filter(lambda x: len(x) > 0, self.img_files[data_name]))
|
|
|
|
|
|
|
|
self.img_start_index[data_name] = img_index
|
|
|
|
img_index += len(self.img_files[data_name])
|
|
|
|
|
|
|
|
# record label_files
|
|
|
|
self.label_files[data_name] = [
|
|
|
|
x.replace('images', 'labels_with_ids').replace(
|
|
|
|
'.png', '.txt').replace('.jpg', '.txt')
|
|
|
|
for x in self.img_files[data_name]
|
|
|
|
]
|
|
|
|
|
|
|
|
for data_name, label_paths in self.label_files.items():
|
|
|
|
max_index = -1
|
|
|
|
for lp in label_paths:
|
|
|
|
lb = np.loadtxt(lp)
|
|
|
|
if len(lb) < 1:
|
|
|
|
continue
|
|
|
|
if len(lb.shape) < 2:
|
|
|
|
img_max = lb[1]
|
|
|
|
else:
|
|
|
|
img_max = np.max(lb[:, 1])
|
|
|
|
if img_max > max_index:
|
|
|
|
max_index = img_max
|
|
|
|
self.tid_num[data_name] = int(max_index + 1)
|
|
|
|
|
|
|
|
last_index = 0
|
|
|
|
for i, (k, v) in enumerate(self.tid_num.items()):
|
|
|
|
self.tid_start_index[k] = last_index
|
|
|
|
last_index += v
|
|
|
|
|
|
|
|
self.num_identities_dict = defaultdict(int)
|
|
|
|
self.num_identities_dict[0] = int(last_index + 1) # single class
|
|
|
|
self.num_imgs_each_data = [len(x) for x in self.img_files.values()]
|
|
|
|
self.total_imgs = sum(self.num_imgs_each_data)
|
|
|
|
|
|
|
|
logger.info('MOT dataset summary: ')
|
|
|
|
logger.info(self.tid_num)
|
|
|
|
logger.info('Total images: {}'.format(self.total_imgs))
|
|
|
|
logger.info('Image start index: {}'.format(self.img_start_index))
|
|
|
|
logger.info('Total identities: {}'.format(self.num_identities_dict[0]))
|
|
|
|
logger.info('Identity start index: {}'.format(self.tid_start_index))
|
|
|
|
|
|
|
|
records = []
|
|
|
|
cname2cid = mot_label()
|
|
|
|
|
|
|
|
for img_index in range(self.total_imgs):
|
|
|
|
for i, (k, v) in enumerate(self.img_start_index.items()):
|
|
|
|
if img_index >= v:
|
|
|
|
data_name = list(self.label_files.keys())[i]
|
|
|
|
start_index = v
|
|
|
|
img_file = self.img_files[data_name][img_index - start_index]
|
|
|
|
lbl_file = self.label_files[data_name][img_index - start_index]
|
|
|
|
|
|
|
|
if not os.path.exists(img_file):
|
|
|
|
logger.warning('Illegal image file: {}, and it will be ignored'.
|
|
|
|
format(img_file))
|
|
|
|
continue
|
|
|
|
if not os.path.isfile(lbl_file):
|
|
|
|
logger.warning('Illegal label file: {}, and it will be ignored'.
|
|
|
|
format(lbl_file))
|
|
|
|
continue
|
|
|
|
|
|
|
|
labels = np.loadtxt(lbl_file, dtype=np.float32).reshape(-1, 6)
|
|
|
|
# each row in labels (N, 6) is [gt_class, gt_identity, cx, cy, w, h]
|
|
|
|
|
|
|
|
cx, cy = labels[:, 2], labels[:, 3]
|
|
|
|
w, h = labels[:, 4], labels[:, 5]
|
|
|
|
gt_bbox = np.stack((cx, cy, w, h)).T.astype('float32')
|
|
|
|
gt_class = labels[:, 0:1].astype('int32')
|
|
|
|
gt_score = np.ones((len(labels), 1)).astype('float32')
|
|
|
|
gt_ide = labels[:, 1:2].astype('int32')
|
|
|
|
for i, _ in enumerate(gt_ide):
|
|
|
|
if gt_ide[i] > -1:
|
|
|
|
gt_ide[i] += self.tid_start_index[data_name]
|
|
|
|
|
|
|
|
mot_rec = {
|
|
|
|
'im_file': img_file,
|
|
|
|
'im_id': img_index,
|
|
|
|
} if 'image' in self.data_fields else {}
|
|
|
|
|
|
|
|
gt_rec = {
|
|
|
|
'gt_class': gt_class,
|
|
|
|
'gt_score': gt_score,
|
|
|
|
'gt_bbox': gt_bbox,
|
|
|
|
'gt_ide': gt_ide,
|
|
|
|
}
|
|
|
|
|
|
|
|
for k, v in gt_rec.items():
|
|
|
|
if k in self.data_fields:
|
|
|
|
mot_rec[k] = v
|
|
|
|
|
|
|
|
records.append(mot_rec)
|
|
|
|
if self.sample_num > 0 and img_index >= self.sample_num:
|
|
|
|
break
|
|
|
|
assert len(records) > 0, 'not found any mot record in %s' % (
|
|
|
|
self.image_lists)
|
|
|
|
self.roidbs, self.cname2cid = records, cname2cid
|
|
|
|
|
|
|
|
|
|
|
|
@register
|
|
|
|
@serializable
|
|
|
|
class MCMOTDataSet(DetDataset):
|
|
|
|
"""
|
|
|
|
Load dataset with MOT format, support multi-class MOT.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
dataset_dir (str): root directory for dataset.
|
|
|
|
image_lists (list(str)): mcmot data image lists, muiti-source mcmot dataset.
|
|
|
|
data_fields (list): key name of data dictionary, at least have 'image'.
|
|
|
|
label_list (str): if use_default_label is False, will load
|
|
|
|
mapping between category and class index.
|
|
|
|
sample_num (int): number of samples to load, -1 means all.
|
|
|
|
|
|
|
|
Notes:
|
|
|
|
MCMOT datasets root directory following this:
|
|
|
|
dataset/mot
|
|
|
|
|——————image_lists
|
|
|
|
| |——————visdrone_mcmot.train
|
|
|
|
| |——————visdrone_mcmot.val
|
|
|
|
visdrone_mcmot
|
|
|
|
|——————images
|
|
|
|
| └——————train
|
|
|
|
| └——————val
|
|
|
|
└——————labels_with_ids
|
|
|
|
└——————train
|
|
|
|
"""
|
|
|
|
|
|
|
|
def __init__(self,
|
|
|
|
dataset_dir=None,
|
|
|
|
image_lists=[],
|
|
|
|
data_fields=['image'],
|
|
|
|
label_list=None,
|
|
|
|
sample_num=-1):
|
|
|
|
super(MCMOTDataSet, self).__init__(
|
|
|
|
dataset_dir=dataset_dir,
|
|
|
|
data_fields=data_fields,
|
|
|
|
sample_num=sample_num)
|
|
|
|
self.dataset_dir = dataset_dir
|
|
|
|
self.image_lists = image_lists
|
|
|
|
if isinstance(self.image_lists, str):
|
|
|
|
self.image_lists = [self.image_lists]
|
|
|
|
self.label_list = label_list
|
|
|
|
self.roidbs = None
|
|
|
|
self.cname2cid = None
|
|
|
|
|
|
|
|
def get_anno(self):
|
|
|
|
if self.image_lists == []:
|
|
|
|
return
|
|
|
|
# only used to get categories and metric
|
|
|
|
# only check first data, but the label_list of all data should be same.
|
|
|
|
first_mot_data = self.image_lists[0].split('.')[0]
|
|
|
|
anno_file = os.path.join(self.dataset_dir, first_mot_data,
|
|
|
|
'label_list.txt')
|
|
|
|
return anno_file
|
|
|
|
|
|
|
|
def parse_dataset(self):
|
|
|
|
self.img_files = OrderedDict()
|
|
|
|
self.img_start_index = OrderedDict()
|
|
|
|
self.label_files = OrderedDict()
|
|
|
|
self.tid_num = OrderedDict()
|
|
|
|
self.tid_start_idx_of_cls_ids = defaultdict(dict) # for MCMOT
|
|
|
|
|
|
|
|
img_index = 0
|
|
|
|
for data_name in self.image_lists:
|
|
|
|
# check every data image list
|
|
|
|
image_lists_dir = os.path.join(self.dataset_dir, 'image_lists')
|
|
|
|
assert os.path.isdir(image_lists_dir), \
|
|
|
|
"The {} is not a directory.".format(image_lists_dir)
|
|
|
|
|
|
|
|
list_path = os.path.join(image_lists_dir, data_name)
|
|
|
|
assert os.path.exists(list_path), \
|
|
|
|
"The list path {} does not exist.".format(list_path)
|
|
|
|
|
|
|
|
# record img_files, filter out empty ones
|
|
|
|
with open(list_path, 'r') as file:
|
|
|
|
self.img_files[data_name] = file.readlines()
|
|
|
|
self.img_files[data_name] = [
|
|
|
|
os.path.join(self.dataset_dir, x.strip())
|
|
|
|
for x in self.img_files[data_name]
|
|
|
|
]
|
|
|
|
self.img_files[data_name] = list(
|
|
|
|
filter(lambda x: len(x) > 0, self.img_files[data_name]))
|
|
|
|
|
|
|
|
self.img_start_index[data_name] = img_index
|
|
|
|
img_index += len(self.img_files[data_name])
|
|
|
|
|
|
|
|
# record label_files
|
|
|
|
self.label_files[data_name] = [
|
|
|
|
x.replace('images', 'labels_with_ids').replace(
|
|
|
|
'.png', '.txt').replace('.jpg', '.txt')
|
|
|
|
for x in self.img_files[data_name]
|
|
|
|
]
|
|
|
|
|
|
|
|
for data_name, label_paths in self.label_files.items():
|
|
|
|
# using max_ids_dict rather than max_index
|
|
|
|
max_ids_dict = defaultdict(int)
|
|
|
|
for lp in label_paths:
|
|
|
|
lb = np.loadtxt(lp)
|
|
|
|
if len(lb) < 1:
|
|
|
|
continue
|
|
|
|
lb = lb.reshape(-1, 6)
|
|
|
|
for item in lb:
|
|
|
|
if item[1] > max_ids_dict[int(item[0])]:
|
|
|
|
# item[0]: cls_id
|
|
|
|
# item[1]: track id
|
|
|
|
max_ids_dict[int(item[0])] = int(item[1])
|
|
|
|
# track id number
|
|
|
|
self.tid_num[data_name] = max_ids_dict
|
|
|
|
|
|
|
|
last_idx_dict = defaultdict(int)
|
|
|
|
for i, (k, v) in enumerate(self.tid_num.items()): # each sub dataset
|
|
|
|
for cls_id, id_num in v.items(): # v is a max_ids_dict
|
|
|
|
self.tid_start_idx_of_cls_ids[k][cls_id] = last_idx_dict[cls_id]
|
|
|
|
last_idx_dict[cls_id] += id_num
|
|
|
|
|
|
|
|
self.num_identities_dict = defaultdict(int)
|
|
|
|
for k, v in last_idx_dict.items():
|
|
|
|
self.num_identities_dict[k] = int(v) # total ids of each category
|
|
|
|
|
|
|
|
self.num_imgs_each_data = [len(x) for x in self.img_files.values()]
|
|
|
|
self.total_imgs = sum(self.num_imgs_each_data)
|
|
|
|
|
|
|
|
# cname2cid and cid2cname
|
|
|
|
cname2cid = {}
|
|
|
|
if self.label_list is not None:
|
|
|
|
# if use label_list for multi source mix dataset,
|
|
|
|
# please make sure label_list in the first sub_dataset at least.
|
|
|
|
sub_dataset = self.image_lists[0].split('.')[0]
|
|
|
|
label_path = os.path.join(self.dataset_dir, sub_dataset,
|
|
|
|
self.label_list)
|
|
|
|
if not os.path.exists(label_path):
|
|
|
|
logger.info(
|
|
|
|
"Note: label_list {} does not exists, use VisDrone 10 classes labels as default.".
|
|
|
|
format(label_path))
|
|
|
|
cname2cid = visdrone_mcmot_label()
|
|
|
|
else:
|
|
|
|
with open(label_path, 'r') as fr:
|
|
|
|
label_id = 0
|
|
|
|
for line in fr.readlines():
|
|
|
|
cname2cid[line.strip()] = label_id
|
|
|
|
label_id += 1
|
|
|
|
else:
|
|
|
|
cname2cid = visdrone_mcmot_label()
|
|
|
|
|
|
|
|
cid2cname = dict([(v, k) for (k, v) in cname2cid.items()])
|
|
|
|
|
|
|
|
logger.info('MCMOT dataset summary: ')
|
|
|
|
logger.info(self.tid_num)
|
|
|
|
logger.info('Total images: {}'.format(self.total_imgs))
|
|
|
|
logger.info('Image start index: {}'.format(self.img_start_index))
|
|
|
|
|
|
|
|
logger.info('Total identities of each category: ')
|
|
|
|
num_identities_dict = sorted(
|
|
|
|
self.num_identities_dict.items(), key=lambda x: x[0])
|
|
|
|
total_IDs_all_cats = 0
|
|
|
|
for (k, v) in num_identities_dict:
|
|
|
|
logger.info('Category {} [{}] has {} IDs.'.format(k, cid2cname[k],
|
|
|
|
v))
|
|
|
|
total_IDs_all_cats += v
|
|
|
|
logger.info('Total identities of all categories: {}'.format(
|
|
|
|
total_IDs_all_cats))
|
|
|
|
|
|
|
|
logger.info('Identity start index of each category: ')
|
|
|
|
for k, v in self.tid_start_idx_of_cls_ids.items():
|
|
|
|
sorted_v = sorted(v.items(), key=lambda x: x[0])
|
|
|
|
for (cls_id, start_idx) in sorted_v:
|
|
|
|
logger.info('Start index of dataset {} category {:d} is {:d}'
|
|
|
|
.format(k, cls_id, start_idx))
|
|
|
|
|
|
|
|
records = []
|
|
|
|
for img_index in range(self.total_imgs):
|
|
|
|
for i, (k, v) in enumerate(self.img_start_index.items()):
|
|
|
|
if img_index >= v:
|
|
|
|
data_name = list(self.label_files.keys())[i]
|
|
|
|
start_index = v
|
|
|
|
img_file = self.img_files[data_name][img_index - start_index]
|
|
|
|
lbl_file = self.label_files[data_name][img_index - start_index]
|
|
|
|
|
|
|
|
if not os.path.exists(img_file):
|
|
|
|
logger.warning('Illegal image file: {}, and it will be ignored'.
|
|
|
|
format(img_file))
|
|
|
|
continue
|
|
|
|
if not os.path.isfile(lbl_file):
|
|
|
|
logger.warning('Illegal label file: {}, and it will be ignored'.
|
|
|
|
format(lbl_file))
|
|
|
|
continue
|
|
|
|
|
|
|
|
labels = np.loadtxt(lbl_file, dtype=np.float32).reshape(-1, 6)
|
|
|
|
# each row in labels (N, 6) is [gt_class, gt_identity, cx, cy, w, h]
|
|
|
|
|
|
|
|
cx, cy = labels[:, 2], labels[:, 3]
|
|
|
|
w, h = labels[:, 4], labels[:, 5]
|
|
|
|
gt_bbox = np.stack((cx, cy, w, h)).T.astype('float32')
|
|
|
|
gt_class = labels[:, 0:1].astype('int32')
|
|
|
|
gt_score = np.ones((len(labels), 1)).astype('float32')
|
|
|
|
gt_ide = labels[:, 1:2].astype('int32')
|
|
|
|
for i, _ in enumerate(gt_ide):
|
|
|
|
if gt_ide[i] > -1:
|
|
|
|
cls_id = int(gt_class[i])
|
|
|
|
start_idx = self.tid_start_idx_of_cls_ids[data_name][cls_id]
|
|
|
|
gt_ide[i] += start_idx
|
|
|
|
|
|
|
|
mot_rec = {
|
|
|
|
'im_file': img_file,
|
|
|
|
'im_id': img_index,
|
|
|
|
} if 'image' in self.data_fields else {}
|
|
|
|
|
|
|
|
gt_rec = {
|
|
|
|
'gt_class': gt_class,
|
|
|
|
'gt_score': gt_score,
|
|
|
|
'gt_bbox': gt_bbox,
|
|
|
|
'gt_ide': gt_ide,
|
|
|
|
}
|
|
|
|
|
|
|
|
for k, v in gt_rec.items():
|
|
|
|
if k in self.data_fields:
|
|
|
|
mot_rec[k] = v
|
|
|
|
|
|
|
|
records.append(mot_rec)
|
|
|
|
if self.sample_num > 0 and img_index >= self.sample_num:
|
|
|
|
break
|
|
|
|
assert len(records) > 0, 'not found any mot record in %s' % (
|
|
|
|
self.image_lists)
|
|
|
|
self.roidbs, self.cname2cid = records, cname2cid
|
|
|
|
|
|
|
|
|
|
|
|
@register
|
|
|
|
@serializable
|
|
|
|
class MOTImageFolder(DetDataset):
|
|
|
|
"""
|
|
|
|
Load MOT dataset with MOT format from image folder or video .
|
|
|
|
Args:
|
|
|
|
video_file (str): path of the video file, default ''.
|
|
|
|
frame_rate (int): frame rate of the video, use cv2 VideoCapture if not set.
|
|
|
|
dataset_dir (str): root directory for dataset.
|
|
|
|
keep_ori_im (bool): whether to keep original image, default False.
|
|
|
|
Set True when used during MOT model inference while saving
|
|
|
|
images or video, or used in DeepSORT.
|
|
|
|
"""
|
|
|
|
|
|
|
|
def __init__(self,
|
|
|
|
video_file=None,
|
|
|
|
frame_rate=-1,
|
|
|
|
dataset_dir=None,
|
|
|
|
data_root=None,
|
|
|
|
image_dir=None,
|
|
|
|
sample_num=-1,
|
|
|
|
keep_ori_im=False,
|
|
|
|
**kwargs):
|
|
|
|
super(MOTImageFolder, self).__init__(
|
|
|
|
dataset_dir, image_dir, sample_num=sample_num)
|
|
|
|
self.video_file = video_file
|
|
|
|
self.data_root = data_root
|
|
|
|
self.keep_ori_im = keep_ori_im
|
|
|
|
self._imid2path = {}
|
|
|
|
self.roidbs = None
|
|
|
|
self.frame_rate = frame_rate
|
|
|
|
|
|
|
|
def check_or_download_dataset(self):
|
|
|
|
return
|
|
|
|
|
|
|
|
def parse_dataset(self, ):
|
|
|
|
if not self.roidbs:
|
|
|
|
if self.video_file is None:
|
|
|
|
self.frame_rate = 30 # set as default if infer image folder
|
|
|
|
self.roidbs = self._load_images()
|
|
|
|
else:
|
|
|
|
self.roidbs = self._load_video_images()
|
|
|
|
|
|
|
|
def _load_video_images(self):
|
|
|
|
if self.frame_rate == -1:
|
|
|
|
# if frame_rate is not set for video, use cv2.VideoCapture
|
|
|
|
cap = cv2.VideoCapture(self.video_file)
|
|
|
|
self.frame_rate = int(cap.get(cv2.CAP_PROP_FPS))
|
|
|
|
|
|
|
|
extension = self.video_file.split('.')[-1]
|
|
|
|
output_path = self.video_file.replace('.{}'.format(extension), '')
|
|
|
|
frames_path = video2frames(self.video_file, output_path,
|
|
|
|
self.frame_rate)
|
|
|
|
self.video_frames = sorted(
|
|
|
|
glob.glob(os.path.join(frames_path, '*.png')))
|
|
|
|
|
|
|
|
self.video_length = len(self.video_frames)
|
|
|
|
logger.info('Length of the video: {:d} frames.'.format(
|
|
|
|
self.video_length))
|
|
|
|
ct = 0
|
|
|
|
records = []
|
|
|
|
for image in self.video_frames:
|
|
|
|
assert image != '' and os.path.isfile(image), \
|
|
|
|
"Image {} not found".format(image)
|
|
|
|
if self.sample_num > 0 and ct >= self.sample_num:
|
|
|
|
break
|
|
|
|
rec = {'im_id': np.array([ct]), 'im_file': image}
|
|
|
|
if self.keep_ori_im:
|
|
|
|
rec.update({'keep_ori_im': 1})
|
|
|
|
self._imid2path[ct] = image
|
|
|
|
ct += 1
|
|
|
|
records.append(rec)
|
|
|
|
assert len(records) > 0, "No image file found"
|
|
|
|
return records
|
|
|
|
|
|
|
|
def _find_images(self):
|
|
|
|
image_dir = self.image_dir
|
|
|
|
if not isinstance(image_dir, Sequence):
|
|
|
|
image_dir = [image_dir]
|
|
|
|
images = []
|
|
|
|
for im_dir in image_dir:
|
|
|
|
if os.path.isdir(im_dir):
|
|
|
|
im_dir = os.path.join(self.dataset_dir, im_dir)
|
|
|
|
images.extend(_make_dataset(im_dir))
|
|
|
|
elif os.path.isfile(im_dir) and _is_valid_file(im_dir):
|
|
|
|
images.append(im_dir)
|
|
|
|
return images
|
|
|
|
|
|
|
|
def _load_images(self):
|
|
|
|
images = self._find_images()
|
|
|
|
ct = 0
|
|
|
|
records = []
|
|
|
|
for image in images:
|
|
|
|
assert image != '' and os.path.isfile(image), \
|
|
|
|
"Image {} not found".format(image)
|
|
|
|
if self.sample_num > 0 and ct >= self.sample_num:
|
|
|
|
break
|
|
|
|
rec = {'im_id': np.array([ct]), 'im_file': image}
|
|
|
|
if self.keep_ori_im:
|
|
|
|
rec.update({'keep_ori_im': 1})
|
|
|
|
self._imid2path[ct] = image
|
|
|
|
ct += 1
|
|
|
|
records.append(rec)
|
|
|
|
assert len(records) > 0, "No image file found"
|
|
|
|
return records
|
|
|
|
|
|
|
|
def get_imid2path(self):
|
|
|
|
return self._imid2path
|
|
|
|
|
|
|
|
def set_images(self, images):
|
|
|
|
self.image_dir = images
|
|
|
|
self.roidbs = self._load_images()
|
|
|
|
|
|
|
|
def set_video(self, video_file, frame_rate):
|
|
|
|
# update video_file and frame_rate by command line of tools/infer_mot.py
|
|
|
|
self.video_file = video_file
|
|
|
|
self.frame_rate = frame_rate
|
|
|
|
assert os.path.isfile(self.video_file) and _is_valid_video(self.video_file), \
|
|
|
|
"wrong or unsupported file format: {}".format(self.video_file)
|
|
|
|
self.roidbs = self._load_video_images()
|
|
|
|
|
|
|
|
|
|
|
|
def _is_valid_video(f, extensions=('.mp4', '.avi', '.mov', '.rmvb', 'flv')):
|
|
|
|
return f.lower().endswith(extensions)
|
|
|
|
|
|
|
|
|
|
|
|
def video2frames(video_path, outpath, frame_rate, **kargs):
|
|
|
|
def _dict2str(kargs):
|
|
|
|
cmd_str = ''
|
|
|
|
for k, v in kargs.items():
|
|
|
|
cmd_str += (' ' + str(k) + ' ' + str(v))
|
|
|
|
return cmd_str
|
|
|
|
|
|
|
|
ffmpeg = ['ffmpeg ', ' -y -loglevel ', ' error ']
|
|
|
|
vid_name = os.path.basename(video_path).split('.')[0]
|
|
|
|
out_full_path = os.path.join(outpath, vid_name)
|
|
|
|
|
|
|
|
if not os.path.exists(out_full_path):
|
|
|
|
os.makedirs(out_full_path)
|
|
|
|
|
|
|
|
# video file name
|
|
|
|
outformat = os.path.join(out_full_path, '%08d.png')
|
|
|
|
|
|
|
|
cmd = ffmpeg
|
|
|
|
cmd = ffmpeg + [
|
|
|
|
' -i ', video_path, ' -r ', str(frame_rate), ' -f image2 ', outformat
|
|
|
|
]
|
|
|
|
cmd = ''.join(cmd) + _dict2str(kargs)
|
|
|
|
|
|
|
|
if os.system(cmd) != 0:
|
|
|
|
raise RuntimeError('ffmpeg process video: {} error'.format(video_path))
|
|
|
|
sys.exit(-1)
|
|
|
|
|
|
|
|
sys.stdout.flush()
|
|
|
|
return out_full_path
|
|
|
|
|
|
|
|
|
|
|
|
def mot_label():
|
|
|
|
labels_map = {'person': 0}
|
|
|
|
return labels_map
|
|
|
|
|
|
|
|
|
|
|
|
def visdrone_mcmot_label():
|
|
|
|
labels_map = {
|
|
|
|
'pedestrian': 0,
|
|
|
|
'people': 1,
|
|
|
|
'bicycle': 2,
|
|
|
|
'car': 3,
|
|
|
|
'van': 4,
|
|
|
|
'truck': 5,
|
|
|
|
'tricycle': 6,
|
|
|
|
'awning-tricycle': 7,
|
|
|
|
'bus': 8,
|
|
|
|
'motor': 9,
|
|
|
|
}
|
|
|
|
return labels_map
|