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")