|
5 | 5 | import tempfile
|
6 | 6 | import unittest
|
7 | 7 |
|
| 8 | +import cv2 |
8 | 9 | import numpy as np
|
9 | 10 | from mmcv import Config
|
10 | 11 | from PIL import Image
|
11 |
| -from tests.ut_config import (MODEL_CONFIG_SEGFORMER, |
| 12 | +from tests.ut_config import (MODEL_CONFIG_MASK2FORMER_INS, |
| 13 | + MODEL_CONFIG_MASK2FORMER_PAN, |
| 14 | + MODEL_CONFIG_SEGFORMER, |
12 | 15 | PRETRAINED_MODEL_MASK2FORMER_DIR,
|
13 | 16 | PRETRAINED_MODEL_SEGFORMER, TEST_IMAGES_DIR)
|
14 | 17 |
|
15 | 18 | from easycv.file import io
|
16 |
| -from easycv.predictors.segmentation import SegmentationPredictor |
| 19 | +from easycv.predictors.segmentation import (Mask2formerPredictor, |
| 20 | + SegmentationPredictor) |
17 | 21 |
|
18 | 22 |
|
19 | 23 | class SegmentationPredictorTest(unittest.TestCase):
|
@@ -112,34 +116,94 @@ def test_dump(self):
|
112 | 116 | shutil.rmtree(temp_dir, ignore_errors=True)
|
113 | 117 |
|
114 | 118 |
|
115 |
| -@unittest.skipIf(True, 'WIP') |
116 | 119 | class Mask2formerPredictorTest(unittest.TestCase):
|
117 | 120 |
|
118 |
| - def test_single(self): |
119 |
| - import cv2 |
120 |
| - from easycv.predictors.segmentation import Mask2formerPredictor |
121 |
| - pan_ckpt = os.path.join(PRETRAINED_MODEL_MASK2FORMER_DIR, |
122 |
| - 'mask2former_pan_export.pth') |
123 |
| - instance_ckpt = os.path.join(PRETRAINED_MODEL_MASK2FORMER_DIR, |
124 |
| - 'mask2former_r50_instance.pth') |
125 |
| - img_path = os.path.join(TEST_IMAGES_DIR, 'mask2former.jpg') |
| 121 | + def setUp(self): |
| 122 | + print(('Testing %s.%s' % (type(self).__name__, self._testMethodName))) |
| 123 | + self.img_path = './data/test/segmentation/data/000000309022.jpg' |
| 124 | + self.pan_ckpt = './data/test/segmentation/models/mask2former_pan_export.pth' |
| 125 | + self.instance_ckpt = './data/test/segmentation/models/mask2former_r50_instance.pth' |
126 | 126 |
|
| 127 | + def test_panoptic_single(self): |
127 | 128 | # panop
|
| 129 | + segmentation_model_config = MODEL_CONFIG_MASK2FORMER_PAN |
128 | 130 | predictor = Mask2formerPredictor(
|
129 |
| - model_path=pan_ckpt, output_mode='panoptic') |
130 |
| - img = cv2.imread(img_path) |
131 |
| - predict_out = predictor([img]) |
132 |
| - pan_img = predictor.show_panoptic(img, predict_out[0]['pan']) |
| 131 | + model_path=self.pan_ckpt, |
| 132 | + task_mode='panoptic', |
| 133 | + config_file=segmentation_model_config) |
| 134 | + img = cv2.imread(self.img_path) |
| 135 | + predict_out = predictor([self.img_path]) |
| 136 | + self.assertEqual(len(predict_out), 1) |
| 137 | + self.assertEqual(len(predict_out[0]['masks']), 14) |
| 138 | + self.assertListEqual( |
| 139 | + predict_out[0]['labels_ids'].tolist(), |
| 140 | + [71, 69, 39, 39, 39, 128, 127, 122, 118, 115, 111, 104, 84, 83]) |
| 141 | + |
| 142 | + pan_img = predictor.show_panoptic( |
| 143 | + img, |
| 144 | + masks=predict_out[0]['masks'], |
| 145 | + labels=predict_out[0]['labels_ids']) |
133 | 146 | cv2.imwrite('pan_out.jpg', pan_img)
|
134 | 147 |
|
| 148 | + def test_panoptic_batch(self): |
| 149 | + total_samples = 2 |
| 150 | + segmentation_model_config = MODEL_CONFIG_MASK2FORMER_PAN |
| 151 | + predictor = Mask2formerPredictor( |
| 152 | + model_path=self.pan_ckpt, |
| 153 | + task_mode='panoptic', |
| 154 | + config_file=segmentation_model_config, |
| 155 | + batch_size=total_samples) |
| 156 | + predict_out = predictor([self.img_path] * total_samples) |
| 157 | + self.assertEqual(len(predict_out), total_samples) |
| 158 | + img = cv2.imread(self.img_path) |
| 159 | + for i in range(total_samples): |
| 160 | + save_name = 'pan_out_batch_%d.jpg' % i |
| 161 | + self.assertEqual(len(predict_out[i]['masks']), 14) |
| 162 | + self.assertListEqual(predict_out[i]['labels_ids'].tolist(), [ |
| 163 | + 71, 69, 39, 39, 39, 128, 127, 122, 118, 115, 111, 104, 84, 83 |
| 164 | + ]) |
| 165 | + pan_img = predictor.show_panoptic( |
| 166 | + img, |
| 167 | + masks=predict_out[i]['masks'], |
| 168 | + labels=predict_out[i]['labels_ids']) |
| 169 | + cv2.imwrite(save_name, pan_img) |
| 170 | + |
| 171 | + def test_instance_single(self): |
135 | 172 | # instance
|
| 173 | + segmentation_model_config = MODEL_CONFIG_MASK2FORMER_INS |
136 | 174 | predictor = Mask2formerPredictor(
|
137 |
| - model_path=instance_ckpt, output_mode='instance') |
138 |
| - img = cv2.imread(img_path) |
139 |
| - predict_out = predictor.predict([img], mode='instance') |
| 175 | + model_path=self.instance_ckpt, |
| 176 | + task_mode='instance', |
| 177 | + config_file=segmentation_model_config) |
| 178 | + img = cv2.imread(self.img_path) |
| 179 | + predict_out = predictor([self.img_path]) |
| 180 | + self.assertEqual(len(predict_out), 1) |
| 181 | + self.assertEqual(predict_out[0]['segms'].shape, (100, 480, 640)) |
| 182 | + self.assertListEqual(predict_out[0]['labels'][:10].tolist(), |
| 183 | + [41, 69, 72, 45, 68, 70, 41, 69, 69, 45]) |
| 184 | + |
140 | 185 | instance_img = predictor.show_instance(img, **predict_out[0])
|
141 | 186 | cv2.imwrite('instance_out.jpg', instance_img)
|
142 | 187 |
|
| 188 | + def test_instance_batch(self): |
| 189 | + total_samples = 2 |
| 190 | + segmentation_model_config = MODEL_CONFIG_MASK2FORMER_INS |
| 191 | + predictor = Mask2formerPredictor( |
| 192 | + model_path=self.instance_ckpt, |
| 193 | + task_mode='instance', |
| 194 | + config_file=segmentation_model_config, |
| 195 | + batch_size=total_samples) |
| 196 | + img = cv2.imread(self.img_path) |
| 197 | + predict_out = predictor([self.img_path] * total_samples) |
| 198 | + self.assertEqual(len(predict_out), total_samples) |
| 199 | + for i in range(total_samples): |
| 200 | + save_name = 'instance_out_batch_%d.jpg' % i |
| 201 | + self.assertEqual(predict_out[i]['segms'].shape, (100, 480, 640)) |
| 202 | + self.assertListEqual(predict_out[0]['labels'][:10].tolist(), |
| 203 | + [41, 69, 72, 45, 68, 70, 41, 69, 69, 45]) |
| 204 | + instance_img = predictor.show_instance(img, **(predict_out[i])) |
| 205 | + cv2.imwrite(save_name, instance_img) |
| 206 | + |
143 | 207 |
|
144 | 208 | if __name__ == '__main__':
|
145 | 209 | unittest.main()
|
0 commit comments