def test_get_dataloader(self):
batch_size = 8
data_loader = data_helper.get_dataloader(data_root='data-bin', mode='train', batch_size=batch_size)
# Test the dataloader
images, targets = next(iter(data_loader))
self.assertEqual(len(images), batch_size)
self.assertIsInstance(images[0], Tensor)
self.assertEqual(len(images[0]), 3)
self.assertEqual(len(targets), batch_size)
self.assertIsInstance(targets[0], Dict)
self.assertIsInstance(targets[0]["image_id"], Tensor)
self.assertIsInstance(targets[0]["boxes"], Tensor)
self.assertIsInstance(targets[0]["labels"], Tensor)
self.assertIsInstance(targets[0]["orig_size"], Tensor)
def test_vanilla_coco_evaluator(self):
# Acquire the images and labels from the coco128 dataset
val_dataloader = data_helper.get_dataloader(data_root='data-bin', mode='val')
coco = data_helper.get_coco_api_from_dataset(val_dataloader.dataset)
coco_evaluator = COCOEvaluator(coco)
# Load model
model = yolov5s(pretrained=True, score_thresh=0.001)
model.eval()
for images, targets in val_dataloader:
preds = model(images)
coco_evaluator.update(preds, targets)
results = coco_evaluator.compute()
self.assertGreater(results['AP'], 41.5)
self.assertGreater(results['AP50'], 62.0)
def test_get_dataloader():
batch_size = 8
data_loader = data_helper.get_dataloader(data_root="data-bin", mode="train", batch_size=batch_size)
# Test the dataloader
images, targets = next(iter(data_loader))
assert len(images) == batch_size
assert isinstance(images[0], Tensor)
assert len(images[0]) == 3
assert len(targets) == batch_size
assert isinstance(targets[0], dict)
assert isinstance(targets[0]["image_id"], Tensor)
assert isinstance(targets[0]["boxes"], Tensor)
assert isinstance(targets[0]["labels"], Tensor)
assert isinstance(targets[0]["orig_size"], Tensor)
def test_vanilla_coco_evaluator(version, map5095, map50):
# Acquire the images and labels from the coco128 dataset
val_dataloader = data_helper.get_dataloader(data_root="data-bin",
mode="val")
coco = data_helper.get_coco_api_from_dataset(val_dataloader.dataset)
coco_evaluator = COCOEvaluator(coco)
# Load model
model = yolov5s(upstream_version=version, pretrained=True)
model = model.eval()
for images, targets in val_dataloader:
preds = model(images)
coco_evaluator.update(preds, targets)
results = coco_evaluator.compute()
assert results["AP"] > map5095
assert results["AP50"] > map50
def test_test_epoch_end(self):
# Acquire the annotation file
data_path = Path('data-bin')
coco128_dirname = 'coco128'
data_helper.prepare_coco128(data_path, dirname=coco128_dirname)
annotation_file = data_path / coco128_dirname / 'annotations' / 'instances_train2017.json'
# Get dataloader to test
val_dataloader = data_helper.get_dataloader(data_root=data_path, mode='val')
# Load model
model = yolov5s(pretrained=True, score_thresh=0.001, annotation_path=annotation_file)
# test step
trainer = pl.Trainer(max_epochs=1)
trainer.test(model, test_dataloaders=val_dataloader)
# test epoch end
results = model.evaluator.compute()
self.assertGreater(results['AP'], 41.5)
self.assertGreater(results['AP50'], 62.0)
def test_test_epoch_end(arch, version, map5095, map50):
# Acquire the annotation file
data_path = Path("data-bin")
coco128_dirname = "coco128"
data_helper.prepare_coco128(data_path, dirname=coco128_dirname)
annotation_file = data_path / coco128_dirname / "annotations" / "instances_train2017.json"
# Get dataloader to test
val_dataloader = data_helper.get_dataloader(data_root=data_path, mode="val")
# Load model
model = DefaultTask(arch=arch, version=version, pretrained=True, annotation_path=annotation_file)
# test step
trainer = pl.Trainer(max_epochs=1)
trainer.test(model, test_dataloaders=val_dataloader)
# test epoch end
results = model.evaluator.compute()
assert results["AP"] > map5095
assert results["AP50"] > map50
def test_train_with_vanilla_module(self):
"""
For issue #86: <https://github.com/zhiqwang/yolov5-rt-stack/issues/86>
"""
# Define the device
device = torch.device('cpu')
train_dataloader = data_helper.get_dataloader(data_root='data-bin', mode='train')
# Sample a pair of images/targets
images, targets = next(iter(train_dataloader))
images = [img.to(device) for img in images]
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
# Define the model
model = yolov5s(num_classes=80)
model.train()
out = model(images, targets)
self.assertIsInstance(out, Dict)
self.assertIsInstance(out["cls_logits"], Tensor)
self.assertIsInstance(out["bbox_regression"], Tensor)
self.assertIsInstance(out["objectness"], Tensor)
