187 lines
8.1 KiB
Python
187 lines
8.1 KiB
Python
"""Extract joints annotations and match with nuScenes ground truths
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
import json
|
|
import logging
|
|
from collections import defaultdict
|
|
import datetime
|
|
import numpy as np
|
|
|
|
from nuscenes.nuscenes import NuScenes
|
|
from nuscenes.utils import splits
|
|
from utils.misc import get_idx_max, append_cluster
|
|
from utils.nuscenes import select_categories
|
|
from utils.camera import project_3d
|
|
from utils.pifpaf import get_input_data, preprocess_pif
|
|
|
|
|
|
class PreprocessNuscenes:
|
|
"""
|
|
Preprocess Nuscenes dataset
|
|
"""
|
|
CAMERAS = ('CAM_FRONT', 'CAM_FRONT_LEFT', 'CAM_FRONT_RIGHT', 'CAM_BACK', 'CAM_BACK_LEFT', 'CAM_BACK_RIGHT')
|
|
dic_jo = {'train': dict(X=[], Y=[], names=[], kps=[], boxes_3d=[], K=[],
|
|
clst=defaultdict(lambda: defaultdict(list))),
|
|
'val': dict(X=[], Y=[], names=[], kps=[], boxes_3d=[], K=[],
|
|
clst=defaultdict(lambda: defaultdict(list))),
|
|
'test': dict(X=[], Y=[], names=[], kps=[], boxes_3d=[], K=[],
|
|
clst=defaultdict(lambda: defaultdict(list)))
|
|
}
|
|
dic_names = defaultdict(lambda: defaultdict(list))
|
|
|
|
def __init__(self, dir_ann, dir_nuscenes, dataset, iou_min=0.3):
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
self.logger = logging.getLogger(__name__)
|
|
|
|
self.iou_min = iou_min
|
|
self.dir_ann = dir_ann
|
|
dir_out = os.path.join('data', 'arrays')
|
|
assert os.path.exists(dir_nuscenes), "Nuscenes directory does not exists"
|
|
assert os.path.exists(self.dir_ann), "The annotations directory does not exists"
|
|
assert os.path.exists(dir_out), "Joints directory does not exists"
|
|
|
|
now = datetime.datetime.now()
|
|
now_time = now.strftime("%Y%m%d-%H%M")[2:]
|
|
self.path_joints = os.path.join(dir_out, 'joints-' + dataset + '-' + now_time + '.json')
|
|
self.path_names = os.path.join(dir_out, 'names-' + dataset + '-' + now_time + '.json')
|
|
|
|
self.nusc, self.scenes, self.split_train, self.split_val = factory(dataset, dir_nuscenes)
|
|
|
|
def run(self):
|
|
"""
|
|
Prepare arrays for training
|
|
"""
|
|
cnt_scenes = 0
|
|
cnt_samples = 0
|
|
cnt_sd = 0
|
|
cnt_ann = 0
|
|
|
|
start = time.time()
|
|
|
|
for ii, scene in enumerate(self.scenes):
|
|
end_scene = time.time()
|
|
current_token = scene['first_sample_token']
|
|
cnt_scenes += 1
|
|
if ii == 0:
|
|
time_left = "Nan"
|
|
else:
|
|
time_left = str((end_scene-start_scene)/60 * (len(self.scenes) - ii))[:4]
|
|
|
|
sys.stdout.write('\r' + 'Elaborating scene {}, remaining time {} minutes'.format(cnt_scenes, time_left) + '\t\n')
|
|
start_scene = time.time()
|
|
if scene['name'] in self.split_train:
|
|
phase = 'train'
|
|
elif scene['name'] in self.split_val:
|
|
phase = 'val'
|
|
else:
|
|
print("phase name not in training or validation split")
|
|
continue
|
|
|
|
while not current_token == "":
|
|
sample_dic = self.nusc.get('sample', current_token)
|
|
cnt_samples += 1
|
|
|
|
# Extract all the sample_data tokens for each sample
|
|
for cam in self.CAMERAS:
|
|
sd_token = sample_dic['data'][cam]
|
|
cnt_sd += 1
|
|
path_im, boxes_obj, kk = self.nusc.get_sample_data(sd_token, box_vis_level=1) # At least one corner
|
|
|
|
# Extract all the annotations of the person
|
|
boxes_gt = []
|
|
dds = []
|
|
boxes_3d = []
|
|
name = os.path.basename(path_im)
|
|
for box_obj in boxes_obj:
|
|
if box_obj.name[:6] != 'animal':
|
|
general_name = box_obj.name.split('.')[0] + '.' + box_obj.name.split('.')[1]
|
|
else:
|
|
general_name = 'animal'
|
|
if general_name in select_categories('all'):
|
|
box = project_3d(box_obj, kk)
|
|
dd = np.linalg.norm(box_obj.center)
|
|
boxes_gt.append(box)
|
|
dds.append(dd)
|
|
box_3d = box_obj.center.tolist() + box_obj.wlh.tolist()
|
|
boxes_3d.append(box_3d)
|
|
self.dic_names[name]['boxes'].append(box)
|
|
self.dic_names[name]['dds'].append(dd)
|
|
self.dic_names[name]['K'] = kk.tolist()
|
|
|
|
# Run IoU with pifpaf detections and save
|
|
path_pif = os.path.join(self.dir_ann, name + '.pifpaf.json')
|
|
exists = os.path.isfile(path_pif)
|
|
|
|
if exists:
|
|
with open(path_pif, 'r') as file:
|
|
annotations = json.load(file)
|
|
|
|
boxes, keypoints = preprocess_pif(annotations, im_size=None)
|
|
(inputs, _), (uv_kps, uv_boxes, _, _) = get_input_data(boxes, keypoints, kk)
|
|
|
|
for ii, box in enumerate(uv_boxes):
|
|
idx_max, iou_max = get_idx_max(box, boxes_gt)
|
|
|
|
if iou_max > self.iou_min:
|
|
|
|
self.dic_jo[phase]['kps'].append(uv_kps[ii])
|
|
self.dic_jo[phase]['X'].append(inputs[ii])
|
|
self.dic_jo[phase]['Y'].append([dds[idx_max]]) # Trick to make it (nn,1)
|
|
self.dic_jo[phase]['names'].append(name) # One image name for each annotation
|
|
self.dic_jo[phase]['boxes_3d'].append(boxes_3d[idx_max])
|
|
self.dic_jo[phase]['K'] = kk.tolist()
|
|
append_cluster(self.dic_jo, phase, inputs[ii], dds[idx_max], uv_kps[ii])
|
|
boxes_gt.pop(idx_max)
|
|
dds.pop(idx_max)
|
|
boxes_3d.pop(idx_max)
|
|
cnt_ann += 1
|
|
sys.stdout.write('\r' + 'Saved annotations {}'
|
|
.format(cnt_ann) + '\t')
|
|
|
|
current_token = sample_dic['next']
|
|
|
|
with open(os.path.join(self.path_joints), 'w') as f:
|
|
json.dump(self.dic_jo, f)
|
|
with open(os.path.join(self.path_names), 'w') as f:
|
|
json.dump(self.dic_names, f)
|
|
end = time.time()
|
|
|
|
print("\nSaved {} annotations for {} samples in {} scenes. Total time: {:.1f} minutes"
|
|
.format(cnt_ann, cnt_samples, cnt_scenes, (end-start)/60))
|
|
print("\nOutput files:\n{}\n{}\n".format(self.path_names, self.path_joints))
|
|
|
|
|
|
def factory(dataset, dir_nuscenes):
|
|
"""Define dataset type and split training and validation"""
|
|
|
|
assert dataset in ['nuscenes', 'nuscenes_mini', 'nuscenes_teaser']
|
|
|
|
if dataset == 'nuscenes':
|
|
nusc = NuScenes(version='v1.0-trainval', dataroot=dir_nuscenes, verbose=True)
|
|
scenes = nusc.scene
|
|
split_scenes = splits.create_splits_scenes()
|
|
split_train, split_val = split_scenes['train'], split_scenes['val']
|
|
|
|
elif dataset == 'nuscenes_mini':
|
|
nusc = NuScenes(version='v1.0-mini', dataroot=dir_nuscenes, verbose=True)
|
|
scenes = nusc.scene
|
|
split_scenes = splits.create_splits_scenes()
|
|
split_train, split_val = split_scenes['train'], split_scenes['val']
|
|
|
|
else:
|
|
nusc = NuScenes(version='v1.0-trainval', dataroot=dir_nuscenes, verbose=True)
|
|
with open("splits/nuscenes_teaser_scenes.txt", "r") as file:
|
|
teaser_scenes = file.read().splitlines()
|
|
scenes = nusc.scene
|
|
scenes = [scene for scene in scenes if scene['token'] in teaser_scenes]
|
|
with open("splits/split_nuscenes_teaser.json", "r") as file:
|
|
dic_split = json.load(file)
|
|
split_train = [scene['name'] for scene in scenes if scene['token'] in dic_split['train']]
|
|
split_val = [scene['name'] for scene in scenes if scene['token'] in dic_split['val']]
|
|
|
|
return nusc, scenes, split_train, split_val
|