def test_run(self):
path = pathlib.Path('test_log')
if not path.exists():
path.mkdir(parents=True)
runner = ObjectDetectionTrainer()
data_path = pathlib.Path(__file__).parent.joinpath('after_path')
data_path.mkdir(parents=True, exist_ok=True)
convert_tfrecord(
data_path,
pathlib.Path(__file__).parent.joinpath('data/mscoco/'))
history = runner.run('yolo', 'mscoco', {'epochs': 1}, {
'batch_size':
2,
'train_image_directory':
pathlib.Path(__file__).parent.joinpath('data/mscoco/train2014'),
'test_image_directory':
pathlib.Path(__file__).parent.joinpath('data/mscoco/val2014'),
'train_data_directory':
data_path.joinpath('train'),
'test_data_directory':
data_path.joinpath('test')
}, path)
ok_(path.joinpath('model').exists())
shutil.rmtree(data_path)
def test_convert_tfrecord(self):
convert_tfrecord(
self.path,
pathlib.Path(__file__).parent.joinpath('data/mscoco/'))
ok_(self.path.joinpath('train/data0.tfrecord').exists())
ok_(self.path.joinpath('test/data0.tfrecord').exists())
dataset = tf.data.TFRecordDataset(
[str(self.path.joinpath('train/data0.tfrecord'))])
def _parse_image_function(example_proto):
return tf.io.parse_single_sequence_example(
example_proto,
context_features=MSCococDatectionDataset.
image_feature_description,
sequence_features=MSCococDatectionDataset.
label_feature_description)
for raw_record in dataset.map(_parse_image_function).take(1):
image_feature, label_feature = raw_record
image = Image.frombytes('RGB', (image_feature['width'].numpy()[0],
image_feature['height'].numpy()[0]),
image_feature['image'].numpy())
image.save('test.png', "PNG")
subprocess.run(('open test.png'), shell=True)
label = label_feature['label'].numpy()
ok_((np.abs(label - [[493.015, 208.61, 214.97, 297.16, 25],
[119.025, 384.085, 132.03, 55.19, 25]]) <=
0.0001).all())
pathlib.Path('test.png').unlink()
def test_draw_resized_bbox(self):
convert_tfrecord(
self.path,
pathlib.Path(__file__).parent.joinpath('data/mscoco/'))
dataset = MSCococDatectionDataset(
train_image_directory=pathlib.Path(__file__).parent.joinpath(
'data/mscoco/train2014'),
test_image_directory=pathlib.Path(__file__).parent.joinpath(
'data/mscoco/val2014'),
train_data_directory=self.path.joinpath('train'),
test_data_directory=self.path.joinpath('test'),
batch_size=2)
image, boxes = next(dataset.training_data_generator())
image = Image.fromarray((image.numpy()[0] * 255.0).astype(np.uint8))
draw = ImageDraw.Draw(image)
for box in boxes[0]:
x, y, w, h, category = box
x *= image.width
w *= image.width
y *= image.height
h *= image.height
draw.rectangle(((x - w / 2, y - h / 2), (x + w / 2, y + h / 2)),
outline='red',
width=5)
image.save('test.png', "PNG")
subprocess.run(('open test.png'), shell=True)
pathlib.Path('test.png').unlink()
def test_init(self):
convert_tfrecord(
self.path,
pathlib.Path(__file__).parent.joinpath('data/mscoco/'))
dataset = MSCococDatectionDataset(
train_image_directory=pathlib.Path(__file__).parent.joinpath(
'data/mscoco/train2014'),
test_image_directory=pathlib.Path(__file__).parent.joinpath(
'data/mscoco/val2014'),
train_data_directory=self.path.joinpath('train'),
test_data_directory=self.path.joinpath('test'),
batch_size=2)
ok_(len(dataset.x_train) != 0)
ok_(isinstance(dataset.x_train[0], pathlib.Path))
def test_eval_generator(self):
convert_tfrecord(
self.path,
pathlib.Path(__file__).parent.joinpath('data/mscoco/'))
dataset = MSCococDatectionDataset(
train_image_directory=pathlib.Path(__file__).parent.joinpath(
'data/mscoco/train2014'),
test_image_directory=pathlib.Path(__file__).parent.joinpath(
'data/mscoco/val2014'),
train_data_directory=self.path.joinpath('train'),
test_data_directory=self.path.joinpath('test'),
batch_size=2)
generator = dataset.eval_data_generator()
image, box = next(generator)
eq_(len(image), 2)
eq_(len(box[0]), dataset.max_boxes)
image, box = next(generator)
eq_(len(image), 2)
def test_preprocess_gt_boxes_with_real_data(self):
path = pathlib.Path(__file__).parent.joinpath('after_path')
path.mkdir(parents=True, exist_ok=True)
convert_tfrecord(
path,
pathlib.Path(__file__).parent.joinpath('data/mscoco/'))
dataset = MSCococDatectionDataset(
train_image_directory=pathlib.Path(__file__).parent.joinpath(
'data/mscoco/train2014'),
test_image_directory=pathlib.Path(__file__).parent.joinpath(
'data/mscoco/val2014'),
train_data_directory=path.joinpath('train'),
test_data_directory=path.joinpath('test'),
batch_size=2)
image, label = next(dataset.training_data_generator())
image = Image.fromarray((image.numpy()[0] * 255.0).astype(np.uint8))
yolo = YoloV2(dataset=dataset)
detect_masks, matching_boxes = yolo.preprocess_gt_boxes(label)
non_zeroes = np.argwhere(detect_masks.numpy()[0][..., 0])
draw = ImageDraw.Draw(image)
for box in non_zeroes:
x, y, w, h, category = matching_boxes.numpy()[0, box[0], box[1],
box[2]]
x = (x + box[1]) * 32
w = np.exp(w) * yolo.anchors[box[2]][0] * 32
y = (y + box[0]) * 32
h = np.exp(h) * yolo.anchors[box[2]][1] * 32
draw.rectangle(((x - w / 2, y - h / 2), (x + w / 2, y + h / 2)),
outline='red',
width=5)
image.save('test.png', "PNG")
subprocess.run(('open test.png'), shell=True)
time.sleep(1)
pathlib.Path('test.png').unlink()
shutil.rmtree(path)