def test_yolov5s_r31(self):
images_one, images_two = self.get_test_images()
images_dummy = [torch.ones(3, 100, 100) * 0.3]
model = yolov5s(upstream_version='v3.1',
export_friendly=True,
pretrained=True)
model.eval()
model(images_one)
# Test exported model on images of different size, or dummy input
self.run_model(model, [(images_one, ), (images_two, ),
(images_dummy, )],
input_names=["images_tensors"],
output_names=["outputs"],
dynamic_axes={
"images_tensors": [0, 1, 2],
"outputs": [0, 1, 2]
},
tolerate_small_mismatch=True)
# Test exported model for an image with no detections on other images
self.run_model(model, [(images_dummy, ), (images_one, )],
input_names=["images_tensors"],
output_names=["outputs"],
dynamic_axes={
"images_tensors": [0, 1, 2],
"outputs": [0, 1, 2]
},
tolerate_small_mismatch=True)
def test_train_one_step(self):
# Load model
model = yolov5s()
model.train()
# Setup the DataModule
train_dataset = DummyCOCODetectionDataset(num_samples=128)
datamodule = DetectionDataModule(train_dataset, batch_size=16)
# Trainer
trainer = pl.Trainer(max_epochs=1)
trainer.fit(model, datamodule)
def test_training_step(self):
# Setup the DataModule
data_path = 'data-bin'
train_dataset = data_helper.get_dataset(data_root=data_path, mode='train')
val_dataset = data_helper.get_dataset(data_root=data_path, mode='val')
data_module = DetectionDataModule(train_dataset, val_dataset, batch_size=16)
# Load model
model = yolov5s()
model.train()
# Trainer
trainer = pl.Trainer(max_epochs=1)
trainer.fit(model, data_module)
def test_predict_with_image_file(self):
# Set image inputs
img_name = "test/assets/zidane.jpg"
# Load model
model = yolov5s(pretrained=True)
model.eval()
# Perform inference on an image file
predictions = model.predict(img_name)
self.assertIsInstance(predictions, list)
self.assertEqual(len(predictions), 1)
self.assertIsInstance(predictions[0], Dict)
self.assertIsInstance(predictions[0]["boxes"], Tensor)
self.assertIsInstance(predictions[0]["labels"], Tensor)
self.assertIsInstance(predictions[0]["scores"], Tensor)
def test_yolov5s_script(self):
model = yolov5s(pretrained=True)
model.eval()
scripted_model = torch.jit.script(model)
scripted_model.eval()
x = [torch.rand(3, 416, 320), torch.rand(3, 480, 352)]
out = model(x)
out_script = scripted_model(x)
self.assertTrue(out[0]["scores"].equal(out_script[0]["scores"]))
self.assertTrue(out[0]["labels"].equal(out_script[0]["labels"]))
self.assertTrue(out[0]["boxes"].equal(out_script[0]["boxes"]))
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_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_predict_with_image_files():
# Set image inputs
img_name1 = "test/assets/zidane.jpg"
img_name2 = "test/assets/bus.jpg"
img_names = [img_name1, img_name2]
# Load model
model = yolov5s(pretrained=True)
model.eval()
# Perform inference on a list of image files
predictions = model.predict(img_names)
assert isinstance(predictions, list)
assert len(predictions) == 2
assert isinstance(predictions[0], dict)
assert isinstance(predictions[0]["boxes"], Tensor)
assert isinstance(predictions[0]["labels"], Tensor)
assert isinstance(predictions[0]["scores"], Tensor)
def test_predict_with_tensor():
# Set image inputs
img_name = "test/assets/zidane.jpg"
img_tensor = default_loader(img_name)
assert img_tensor.ndim == 3
# Load model
model = yolov5s(pretrained=True)
model.eval()
# Perform inference on a list of image files
predictions = model.predict(img_tensor)
assert isinstance(predictions, list)
assert len(predictions) == 1
assert isinstance(predictions[0], dict)
assert isinstance(predictions[0]["boxes"], Tensor)
assert isinstance(predictions[0]["labels"], Tensor)
assert isinstance(predictions[0]["scores"], Tensor)
def test_predict_tensor(self):
# Set image inputs
img_name = "test/assets/zidane.jpg"
img_tensor = default_loader(img_name)
self.assertEqual(img_tensor.ndim, 3)
# Load model
model = yolov5s(pretrained=True)
model.eval()
# Perform inference on a list of image files
predictions = model.predict(img_tensor)
self.assertIsInstance(predictions, list)
self.assertEqual(len(predictions), 1)
self.assertIsInstance(predictions[0], Dict)
self.assertIsInstance(predictions[0]["boxes"], torch.Tensor)
self.assertIsInstance(predictions[0]["labels"], torch.Tensor)
self.assertIsInstance(predictions[0]["scores"], torch.Tensor)
def test_predict_image_files(self):
# Set image inputs
img_name1 = "test/assets/zidane.jpg"
img_name2 = "test/assets/bus.jpg"
img_names = [img_name1, img_name2]
# Load model
model = yolov5s(pretrained=True)
model.eval()
# Perform inference on a list of image files
predictions = model.predict(img_names)
self.assertIsInstance(predictions, list)
self.assertEqual(len(predictions), 2)
self.assertIsInstance(predictions[0], Dict)
self.assertIsInstance(predictions[0]["boxes"], torch.Tensor)
self.assertIsInstance(predictions[0]["labels"], torch.Tensor)
self.assertIsInstance(predictions[0]["scores"], torch.Tensor)
def test_predict_with_vanilla_model(self):
# Set image inputs
img_name = "test/assets/zidane.jpg"
img_input = default_loader(img_name)
self.assertEqual(img_input.ndim, 3)
# Load model
model = yolov5s(pretrained=True)
model.eval()
# Perform inference on a list of tensors
out = model([img_input])
self.assertIsInstance(out, list)
self.assertEqual(len(out), 1)
self.assertIsInstance(out[0], Dict)
self.assertIsInstance(out[0]["boxes"], Tensor)
self.assertIsInstance(out[0]["labels"], Tensor)
self.assertIsInstance(out[0]["scores"], Tensor)
def test_predict_with_vanilla_model():
# Set image inputs
img_name = "test/assets/zidane.jpg"
img_input = default_loader(img_name)
assert img_input.ndim == 3
# Load model
model = yolov5s(pretrained=True)
model.eval()
# Perform inference on a list of tensors
out = model([img_input])
assert isinstance(out, list)
assert len(out) == 1
assert isinstance(out[0], dict)
assert isinstance(out[0]["boxes"], Tensor)
assert isinstance(out[0]["labels"], Tensor)
assert isinstance(out[0]["scores"], Tensor)
def test_predict_with_tensors():
# Set image inputs
img_tensor1 = default_loader("test/assets/zidane.jpg")
assert img_tensor1.ndim == 3
img_tensor2 = default_loader("test/assets/bus.jpg")
assert img_tensor2.ndim == 3
img_tensors = [img_tensor1, img_tensor2]
# Load model
model = yolov5s(pretrained=True)
model.eval()
# Perform inference on a list of image files
predictions = model.predict(img_tensors)
assert isinstance(predictions, list)
assert len(predictions) == 2
assert isinstance(predictions[0], dict)
assert isinstance(predictions[0]["boxes"], Tensor)
assert isinstance(predictions[0]["labels"], Tensor)
assert isinstance(predictions[0]["scores"], Tensor)
def test_predict_with_tensors(self):
# Set image inputs
img_tensor1 = default_loader("test/assets/zidane.jpg")
self.assertEqual(img_tensor1.ndim, 3)
img_tensor2 = default_loader("test/assets/bus.jpg")
self.assertEqual(img_tensor2.ndim, 3)
img_tensors = [img_tensor1, img_tensor2]
# Load model
model = yolov5s(pretrained=True)
model.eval()
# Perform inference on a list of image files
predictions = model.predict(img_tensors)
self.assertIsInstance(predictions, list)
self.assertEqual(len(predictions), 2)
self.assertIsInstance(predictions[0], Dict)
self.assertIsInstance(predictions[0]["boxes"], Tensor)
self.assertIsInstance(predictions[0]["labels"], Tensor)
self.assertIsInstance(predictions[0]["scores"], Tensor)
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_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)
def test_get_trace_module(h, w):
model_func = yolov5s(pretrained=True)
script_module = get_trace_module(model_func, input_shape=(h, w))
assert isinstance(script_module, TopLevelTracedModule)
assert script_module.code is not None
from yolort.models import yolov5s
from yolort.relay import get_trace_module
if __name__ == "__main__":
model = yolov5s(pretrained=True)
traced_model = get_trace_module(model, input_shape=(416, 352))
traced_model.save("./test/tracing/yolov5s.torchscript.pt")
