-
Notifications
You must be signed in to change notification settings - Fork 152
/
test_engine.py
150 lines (130 loc) · 5.65 KB
/
test_engine.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
import unittest
import torch
import pytorch_lightning as pl
from yolort.models.yolo import yolov5_darknet_pan_s_r31
from yolort.models.transform import nested_tensor_from_tensor_list
from yolort.models import yolov5s
from yolort.datasets import DetectionDataModule
from .dataset_utils import DummyCOCODetectionDataset
from typing import Dict
from torchvision.io import read_image
def default_loader(img_name, is_half=False):
"""
Read Image using TorchVision.io Here
"""
img = read_image(img_name)
img = img.half() if is_half else img.float() # uint8 to fp16/32
img /= 255. # 0 - 255 to 0.0 - 1.0
return img
class EngineTester(unittest.TestCase):
def test_train(self):
# Do forward over image
img_name = "test/assets/zidane.jpg"
img_tensor = default_loader(img_name)
self.assertEqual(img_tensor.ndim, 3)
# Add a dummy image to train
img_dummy = torch.rand((3, 416, 360), dtype=torch.float32)
images = nested_tensor_from_tensor_list([img_tensor, img_dummy])
targets = torch.tensor([[0, 7, 0.3790, 0.5487, 0.3220, 0.2047],
[0, 2, 0.2680, 0.5386, 0.2200, 0.1779],
[0, 3, 0.1720, 0.5403, 0.1960, 0.1409],
[0, 4, 0.2240, 0.4547, 0.1520, 0.0705]], dtype=torch.float)
model = yolov5_darknet_pan_s_r31(num_classes=12)
model.train()
out = model(images, targets)
self.assertIsInstance(out, Dict)
self.assertIsInstance(out["cls_logits"], torch.Tensor)
self.assertIsInstance(out["bbox_regression"], torch.Tensor)
self.assertIsInstance(out["objectness"], torch.Tensor)
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_inference(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"], torch.Tensor)
self.assertIsInstance(out[0]["labels"], torch.Tensor)
self.assertIsInstance(out[0]["scores"], torch.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_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"], torch.Tensor)
self.assertIsInstance(predictions[0]["labels"], torch.Tensor)
self.assertIsInstance(predictions[0]["scores"], torch.Tensor)
def test_predict_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"], 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)
if __name__ == '__main__':
unittest.main()