Correctly access one training image

Signed-off-by: Jim Martens <github@2martens.de>
This commit is contained in:
2019-06-24 11:04:57 +02:00
parent 25c4908825
commit 180556f6ba

View File

@ -86,9 +86,9 @@ def _ssd_train(args: argparse.Namespace) -> None:
if args.debug: if args.debug:
train_data = next(train_generator) train_data = next(train_generator)
train_image = train_data[0] train_image = train_data[0][0]
print(train_image) print(train_image)
train_length -= 1 train_length -= batch_size
from PIL import Image from PIL import Image
image = Image.fromarray(train_image) image = Image.fromarray(train_image)
image.save(f"{args.summary_path}/train/{args.network}/{args.iteration}/train_image.png") image.save(f"{args.summary_path}/train/{args.network}/{args.iteration}/train_image.png")