def create_dataset2():
"""
Concat the original images with output of the JNET to train the network
:return:
"""
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)
images2 = loader.load_cached_images(name='jnet_train-200-300.npy',
vi_name='uad_train_vi.npy')
boxes2 = loader.load_cached_boxes(name='uad_train_boxes.npy')
labels2 = loader.load_cached_labels(name='uad_train_labels.npy')
labels2 = labels2['vehicle']
images2, labels2, boxes2 = loader.filter_input3(images2, labels2, boxes2)
final_images = np.concatenate((images, images2), axis=0)
final_boxes = np.concatenate((boxes, boxes2), axis=0)
final_labels = np.concatenate((labels, labels2), axis=0)
logger.info(f"original: {images.shape}, {len(labels)}, {len(boxes)}")
logger.info(
f"final: {final_images.shape}, {final_boxes.shape}, {final_labels.shape}"
)
dataset.set_images(images)
dataset.set_labels(labels)
dataset.set_boxes(boxes)
return dataset, cfg
# load net
num_classes = len(labelmap) + 1 # +1 for background
net = build_ssd('test', 300, num_classes) # initialize SSD
net.load_state_dict(torch.load(args.trained_model))
net.eval()
print('Finished loading 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'))
labels = labels['vehicle']
images, labels, boxes = loader.filter_input3(images, labels, boxes)
dataset = UADDetection(transform=BaseTransform(300, dataset_mean),
target_transform=UADAnnotationTransform())
dataset.set_images(images)
dataset.set_labels(labels)
dataset.set_boxes(boxes)
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# evaluation
"""
Steps:
1. detect all the boxes
2. Save the results in ssd300_uad/test/detections.pkl
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
