new_list = potential_frame_indices.copy()
for filter_i, value in enumerate(potential_frame_indices):
if value == 1:
new_list[filter_i] = final_potential_frame_indices[udf_i]
udf_i += 1
return new_list
if __name__ == "__main__":
### Let's assume you are interested in running a query such as `select FRAMES where FRAME contains VAN`
### load the UADetrac dataset
loaders = UADetracLoader()
images = loaders.load_cached_images()
labels = loaders.load_cached_labels()
boxes = loaders.load_cached_boxes()
### generate binary labels for filters
utils_eva = Utils()
binary_labels = utils_eva.labels2binaryUAD(labels, ['van']) ## can supply as ['car', 'van', 'bus', 'others'] to get binary labels for all
### train the filter -- later we can look into options that save these models
##### https://scikit-learn.org/stable/modules/model_persistence.html
filters = FilterMinimum()
filters.train(images, binary_labels)
### train the UDF
### NOTE: to train the UDF (SSD), it is much more convenient to run `python train_ssd_uad.py` -- I organized the code so that it runs out of the box and trains an SSD model that can be used for evaluation
### Hence, let's assume you ran this code and have a trained SSD model
net.load_state_dict(torch.load(args.trained_model))
net.eval()
logger.info(f"Loaded model {args.trained_model}")
# load data
loader = UADetracLoader()
#images = loader.load_images(dir = os.path.join(home_dir, 'data', 'ua_detrac', '5_images'))
#labels, boxes = loader.load_labels(dir = os.path.join(home_dir, 'data', 'ua_detrac', '5_xml'))
### uad test
#images = loader.load_cached_images(name = 'uad_test_images.npy', vi_name = 'uad_test_vi.npy')
#labels = loader.load_cached_labels(name = 'uad_test_labels.npy')
#boxes = loader.load_cached_boxes(name = 'uad_test_boxes.npy')
### jnet_test
images = loader.load_cached_images(name='jnet_test-200-300.npy',
vi_name='uad_test_vi.npy')
labels = loader.load_cached_labels(name='uad_test_labels.npy')
boxes = loader.load_cached_boxes(name='uad_test_boxes.npy')
labels = labels['vehicle']
images, labels, boxes = loader.filter_input3(images, labels, boxes)
## let's sample images
images = images[::10]
labels = labels[::10]
boxes = boxes[::10]
dataset = UADDetection(transform=BaseTransform(300, dataset_mean),
target_transform=UADAnnotationTransform())
dataset.set_images(images)
new_boxes_per_frame.append(
(left / image_width, top / image_height, right / image_width,
bottom / image_height))
new_boxes.append(new_boxes_per_frame)
assert (len(new_boxes) == len(all_boxes))
for i, boxes_per_frame in enumerate(all_boxes):
assert (len(boxes_per_frame) == len(new_boxes[i]))
return new_boxes
if __name__ == "__main__":
loader = UADetracLoader()
logger = Logger()
images_train = loader.load_cached_images()
labels_train = loader.load_cached_labels()
boxes_train = loader.load_cached_boxes()
### we need to take out frames that don't have labels / boxes
labels_train = labels_train['vehicle']
images_train, labels_train, boxes_train = filter_input(
images_train, labels_train, boxes_train)
image_width, image_height = images_train.shape[1], images_train.shape[2]
logger.info("Finished loading the UADetrac Dataset")
## we need to create a UA Dataset_lite..
## divide the set into train and validation
val_division = int(0.8 * len(images_train))
images_val = images_train[val_division:]
def create_dataset(dataset_name, is_train=None, cache_name=None):
if dataset_name not in available_datasets.keys():
logger.error(
f"dataset: {dataset_name} not in {available_datasets.keys()}")
if (dataset_name is 'UAD'
or dataset_name is 'JNET') and (is_train is None):
logger.error(f"Must specify training or testing for UAD and JNET!")
if dataset_name == 'COCO':
if not os.path.exists(COCO_ROOT):
logger.error('Must specify dataset_root if specifying dataset')
cfg = coco
dataset = COCODetection(root=COCO_ROOT,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
elif dataset_name == 'VOC':
cfg = voc
dataset = VOCDetection(root=VOC_ROOT,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
elif dataset_name == 'UAD' and is_train:
logger.info("We are loading UADetrac!!")
cfg = uad
dataset = UADDetection(
transform=SSDAugmentation(cfg['min_dim'], MEANS))
loader = UADetracLoader()
images = loader.load_cached_images(name='uad_train_images.npy',
vi_name='uad_train_vi.npy')
boxes = loader.load_cached_boxes(name='uad_train_boxes.npy')
labels = loader.load_cached_labels(name='uad_train_labels.npy')
labels = labels['vehicle']
images, labels, boxes = loader.filter_input3(images, labels, boxes)
dataset.set_images(images)
dataset.set_labels(labels)
dataset.set_boxes(boxes)
elif dataset_name == 'UAD' and not is_train:
logger.info("We are loading UADetrac!!")
cfg = uad
dataset = UADDetection(
transform=SSDAugmentation(cfg['min_dim'], MEANS))
loader = UADetracLoader()
images = loader.load_cached_images(name='uad_test_images.npy',
vi_name='uad_test_vi.npy')
boxes = loader.load_cached_boxes(name='uad_test_boxes.npy')
labels = loader.load_cached_labels(name='uad_test_labels.npy')
labels = labels['vehicle']
images, labels, boxes = loader.filter_input3(images, labels, boxes)
images = images[:4000]
labels = labels[:4000]
boxes = boxes[:4000]
dataset.set_images(images)
dataset.set_labels(labels)
dataset.set_boxes(boxes)
elif dataset_name == 'JNET' and is_train:
if cache_name is None:
logger.error("Cache name is required for JNET!! returning...")
return None
logger.info("We are loading JNET - UADetrac!!")
cfg = uad
dataset = UADDetection(
transform=SSDAugmentation(cfg['min_dim'], MEANS))
loader = UADetracLoader()
images = loader.load_cached_images(name=cache_name,
vi_name='uad_train_vi.npy')
labels = loader.load_cached_labels(name='uad_train_labels.npy')
boxes = loader.load_cached_boxes(name='uad_train_boxes.npy')
labels = labels['vehicle']
images, labels, boxes = loader.filter_input3(images, labels, boxes)
logger.info(f"images shape is {images.shape}")
logger.info(f"labels length is {len(labels)}")
logger.info(f"boxes length is {len(boxes)}")
assert (images.shape[0] == len(labels))
assert (len(labels) == len(boxes))
dataset.set_images(images)
dataset.set_labels(labels)
dataset.set_boxes(boxes)
elif dataset_name == 'JNET' and not is_train:
if cache_name is None:
logger.error("Cache name is required for JNET! returning....")
return
logger.info("We are loading JNET - UADetrac!!")
cfg = uad
dataset = UADDetection(
transform=SSDAugmentation(cfg['min_dim'], MEANS))
loader = UADetracLoader()
images = loader.load_cached_images(name=cache_name,
vi_name='uad_test_vi.npy')
labels = loader.load_cached_labels(name='uad_test_labels.npy')
boxes = loader.load_cached_boxes(name='uad_test_boxes.npy')
labels = labels['vehicle']
images, labels, boxes = loader.filter_input3(images, labels, boxes)
###FIXED: we will make this really small so that numbers appear fast
images = images[:2000]
labels = labels[:2000]
boxes = boxes[:2000]
logger.info(f"images shape is {images.shape}")
logger.info(f"labels length is {len(labels)}")
logger.info(f"boxes length is {len(boxes)}")
dataset.set_images(images)
dataset.set_labels(labels)
dataset.set_boxes(boxes)
return dataset, cfg