def build_dataset(cfg):
data_cfg = cfg.DATA_CFG
splits = data_cfg.SPLITS
for split in splits:
dataset_name, standard_format_json_file = get_dataset_name_and_json(data_cfg, split)
# if not os.path.exists(standard_format_json_file) or data_cfg.RECOMPUTE_DATA:
register_dataset(data_cfg, split)
if data_cfg.REPROCESS_RAW_VIDEOS:
raw_to_detectron = RAW_PROCESSORS_MAP[data_cfg.BASE_NAME]
raw_to_detectron(data_cfg, split, standard_format_json_file)
write_coco_format_json(data_cfg, split)
if data_cfg.TRAINED_DETECTOR.DO_INFERENCE:
write_with_inferred_boxes(data_cfg, split)
if data_cfg.TRAINED_DERENDER.DO_INFERENCE:
for attributes_key in data_cfg.TRAINED_DERENDER.ATTRIBUTES_WEIGHTS_MAP:
write_with_inferred_attributes(data_cfg, split, attributes_key)
if data_cfg.SHAPESWORLD_JSON.REPROCESS:
for attributes_to_use in data_cfg.ATTRIBUTES_KEYS:
JsonGenerator(data_cfg, split, "shapesworld", attributes_to_use,
vel_data_assoc="None")
if data_cfg.ADEPT_JSON.REPROCESS:
for attributes_to_use in data_cfg.ADEPT_JSON.ATTRIBUTES_KEYS:
JsonGenerator(data_cfg, split, "adept", attributes_to_use,
vel_data_assoc=data_cfg.ADEPT_JSON.VEL_DATA_ASSOC)
def __init__(self, cfg, split, max_obj=10):
self.cfg = cfg
self.split = split
data_cfg = cfg.DATA_CFG
num_frames = utils.get_num_frames(data_cfg, split)
dataset_name, standard_format_json_file = utils.get_dataset_name_and_json(
data_cfg, split)
print("reading datasets {}".format(dataset_name))
start = time.time()
self.dataset_dicts = utils.get_data_dicts(standard_format_json_file,
num_frames)
# Redoing derender_dataset from trainable_derenderer without detectron
required_fields = [
"pred_box"
] if cfg.MODULE_CFG.DATASETS.USE_PREDICTED_BOXES else ["bbox"]
required_fields += ["attributes"]
# _, self.dataset_dicts = image_based_to_annotation_based(self.dataset_dicts,required_fields)
# I've edited DerenderAttrbutes to basically do nothing. This allows us
# to call it without registering the dataset with detectron, but it also
# loses a lot of the functionality it had.
self.attributes = DerenderAttributes(cfg.MODULE_CFG)
self.mapper = trainers.trainable_derender.ImageBasedDerenderMapper(
cfg.MODULE_CFG.DATASETS.USE_PREDICTED_BOXES,
self.attributes,
False, #for_inference,
use_depth=cfg.MODULE_CFG.DATASETS.USE_DEPTH)
non_visibles = []
for i in range(len(self.dataset_dicts)):
visibles = 0
has_wall = False
has_occluder = False
d = self.dataset_dicts[i]
for annotation in d["annotations"]:
visibles += (annotation["attributes"]["visible"] *
(annotation["attributes"]["type"] == 0)
) #or annotation["attributes"]["type"] == 1))
has_wall = has_wall or (annotation["attributes"]["visible"] and
annotation["attributes"]["type"] == 3)
has_occluder = has_occluder or (
annotation["attributes"]["visible"]
and annotation["attributes"]["type"] == 1)
if (visibles == 0) or has_wall or has_occluder:
non_visibles.append(i)
# non_visibles.reverse()
for idx in range(len(non_visibles) - 1, -1, -1):
self.dataset_dicts.pop(non_visibles[idx])
# self.dataset_dicts = list(map(self.mapper, self.dataset_dicts))
# self.dataset_dicts = [self.mapper(x) for x in self.dataset_dicts]
print("done after {}".format(time.time() - start))
def write_coco_format_json(cfg, split):
timer = CodeTimer("writting to coco")
dataset_name, standard_format_json_file = get_dataset_name_and_json(cfg, split)
dataset_dicts = DatasetCatalog.get(dataset_name)
_,filtered_dicts = filter_dataset(dataset_dicts, required_fields=["bbox", "bbox_mode", "segmentation"])
register_dataset(cfg, split, getter= lambda: filtered_dicts, name=dataset_name+"_for_coco")
coco_dict = convert_to_coco_dict(dataset_name+"_for_coco")
json_format_file = standard_format_json_file.replace(".json", "_coco_format.json")
with open(json_format_file, "w") as f:
json.dump(coco_dict, f)
timer.done()
def __init__(self, cfg, split):
self.cfg = cfg
self.split = split
data_cfg = cfg.DATA_CFG
num_frames = utils.get_num_frames(data_cfg, split)
dataset_name, standard_format_json_file = utils.get_dataset_name_and_json(data_cfg,split)
print("reading datasets {}".format(dataset_name))
start = time.time()
self.dataset_dicts = utils.get_data_dicts(standard_format_json_file, num_frames)
# Redoing derender_dataset from trainable_derenderer without detectron
required_fields = ["pred_box"] if cfg.MODULE_CFG.DATASETS.USE_PREDICTED_BOXES else ["bbox"]
required_fields += ["attributes"]
_, self.dataset_dicts = image_based_to_annotation_based(self.dataset_dicts,required_fields)
def write_with_inferred_boxes(cfg, split):
# TODO: now there are invisible objects the detection mapper ignores that, will have to debug tomorrow
timer = CodeTimer("adding inferred boxes")
module_cfg = os.path.join(cfg.TRAINED_DETECTOR.EXP_DIR, "config.yaml")
module_cfg = load_cfg_from_file(module_cfg)
module_cfg.MODEL.WEIGHTS = cfg.TRAINED_DETECTOR.WEIGHTS_FILE
if cfg.DEBUG:
module_cfg.DATALOADER.NUM_WORKERS = 0
predictor = DetectorPredictor(module_cfg)
dataset_name, standard_format_json_file = get_dataset_name_and_json(
cfg, split)
data_loader = inference_detection_loader(
module_cfg.clone(), dataset_name, DetectionMapper(module_cfg.clone()))
worker_args = []
with torch.no_grad():
for inputs in data_loader:
outputs = predictor(inputs)
for i in range(len(outputs)):
worker_args.append(parse_worker_args(inputs[i], outputs[i]))
if cfg.DEBUG:
new_dicts = [add_inferred_boxes(*w) for w in worker_args]
else:
with Pool(int(cpu_count() / 4)) as p:
new_dicts = p.starmap(add_inferred_boxes, worker_args)
if 'PRED_BOX_SCORE_THRESHOLD' not in cfg:
assert '_val' in split, "start with validation split to compute detection threshold"
cfg.PRED_BOX_SCORE_THRESHOLD = infer_score_threshold(new_dicts)
new_dicts = filter_predicted_boxes_threshold(new_dicts,
cfg.PRED_BOX_SCORE_THRESHOLD)
with open(standard_format_json_file, 'w') as f:
json.dump(new_dicts, f, indent=4)
timer.done()
def build_physics_jsons(self, data_cfg, split):
dataset_name, standard_format_json_file = get_dataset_name_and_json(
data_cfg, split)
dataset = DatasetCatalog.get(dataset_name)
required_fields_values = {self.attributes_key: {"visible": 1}}
_, dataset = filter_dataset(
dataset, required_fields_values=required_fields_values)
videos_dicts = frames2videos(dataset)
out_dir = get_jsons_directory(data_cfg, self.target_physics,
self.attributes_key, dataset_name)
os.makedirs(out_dir, exist_ok=True)
worker_args = [(data_cfg, vid_dict, out_dir, vid_num)
for vid_num, vid_dict in videos_dicts.items()]
if data_cfg.DEBUG:
[self.video2json(*w) for w in worker_args]
else:
with Pool(int(cpu_count())) as p:
p.starmap(self.video2json, worker_args)
self.timer.done()
def write_with_inferred_attributes(cfg, split, attributes_key):
timer = CodeTimer(
"adding inferred attributes split:{}, attributes_key:{}".format(
split, attributes_key))
module_cfg = os.path.join(cfg.TRAINED_DERENDER.EXP_DIR, "cfg.yaml")
module_cfg = load_cfg_from_file(module_cfg)
module_cfg.MODEL.WEIGHTS = cfg.TRAINED_DERENDER.ATTRIBUTES_WEIGHTS_MAP[
attributes_key]
module_cfg.DATALOADER.OBJECTS_PER_BATCH = 1000 if cfg.BASE_NAME == "intphys" else 450
module_cfg.DATALOADER.NUM_WORKERS = 8 if cfg.BASE_NAME == "adept" else module_cfg.DATALOADER.NUM_WORKERS
if cfg.DEBUG:
module_cfg.DATALOADER.NUM_WORKERS = 0
module_cfg.DEBUG = True
module_cfg.DATALOADER.OBJECTS_PER_BATCH = 50
predictor = DerenderPredictor(module_cfg)
# if not cfg.DEBUG:
# gpu_ids = [_ for _ in range(torch.cuda.device_count())]
# predictor.derenderer = torch.nn.parallel.DataParallel(predictor.derenderer, gpu_ids)
dataset_name, standard_format_json_file = get_dataset_name_and_json(
cfg, split)
dataset = DatasetCatalog.get(dataset_name)
required_fields = [
"pred_box"
] if cfg.TRAINED_DERENDER.USE_INFERRED_BOXES else ["bbox"]
filtered_idx, \
mapped_dataset = image_based_to_annotation_based(dataset, required_fields)
mapped_dataset = DatasetFromList(mapped_dataset, copy=False)
mapper = DerenderMapper(cfg.TRAINED_DERENDER.USE_INFERRED_BOXES,
predictor.attributes,
for_inference=True,
use_depth=cfg.TRAINED_DERENDER.USE_DEPTH)
mapped_dataset = MapDataset(mapped_dataset, mapper)
data_loader = DataLoader(
dataset=mapped_dataset,
batch_size=module_cfg.DATALOADER.OBJECTS_PER_BATCH,
num_workers=module_cfg.DATALOADER.NUM_WORKERS,
shuffle=False)
fil_pointer = 0
with torch.no_grad():
for inputs in data_loader:
inputs = to_cuda(inputs)
outputs = predictor(inputs)
batch_size = list(outputs.values())[0].shape[0]
for oix, (img_idx, an_idx) in zip(
range(batch_size),
filtered_idx[fil_pointer:fil_pointer + batch_size]):
dataset[img_idx]["annotations"][an_idx][attributes_key] = \
{k: v[oix].item() for k, v in outputs.items()}
# {k: v[oix].item() if v[oix].size == 1
# else [float(el) for el in v[oix]]
# for k,v in outputs.items()}
fil_pointer = fil_pointer + batch_size
dataset = [fix_for_serialization(d) for d in dataset]
with open(standard_format_json_file, "w") as f:
json.dump(dataset, f, indent=4)
timer.done()