Made it possible to save images without preventing training

Signed-off-by: Jim Martens <github@2martens.de>
This commit is contained in:
2019-07-10 16:29:39 +02:00
parent d517c343be
commit 4c4c05d383
2 changed files with 50 additions and 33 deletions

View File

@ -205,24 +205,23 @@ def _ssd_train(args: argparse.Namespace) -> None:
top_k, top_k,
pre_trained_weights_file) pre_trained_weights_file)
train_generator, train_length, val_generator, val_length = _ssd_train_get_generators( train_generator, train_length, train_debug_generator, \
data.load_scenenet_data, val_generator, val_length, val_debug_generator = _ssd_train_get_generators(args,
file_names_train, data.load_scenenet_data,
instances_train, file_names_train,
file_names_val, instances_train,
instances_val, file_names_val,
coco_path, instances_val,
batch_size, coco_path,
image_size, batch_size,
nr_trajectories, image_size,
predictor_sizes nr_trajectories,
) predictor_sizes)
train_length = _ssd_debug_save_images(args, save_train_images, _ssd_debug_save_images(args, save_train_images,
debug.save_ssd_train_images, coco_utils.get_coco_category_maps, debug.save_ssd_train_images, coco_utils.get_coco_category_maps,
summary_path, coco_path, summary_path, coco_path,
batch_size, image_size, image_size, train_debug_generator)
train_generator, train_length)
nr_batches_train = _get_nr_batches(train_length, batch_size) nr_batches_train = _get_nr_batches(train_length, batch_size)
tensorboard_callback = _ssd_get_tensorboard_callback(args, save_summaries, summary_path) tensorboard_callback = _ssd_get_tensorboard_callback(args, save_summaries, summary_path)
@ -331,7 +330,8 @@ def _ssd_train_load_gt(train_gt_path: str, val_gt_path: str
return file_names_train, instances_train, file_names_val, instances_val return file_names_train, instances_train, file_names_val, instances_val
def _ssd_train_get_generators(load_data: callable, def _ssd_train_get_generators(args: argparse.Namespace,
load_data: callable,
file_names_train: Sequence[Sequence[str]], file_names_train: Sequence[Sequence[str]],
instances_train: Sequence[Sequence[Sequence[dict]]], instances_train: Sequence[Sequence[Sequence[dict]]],
file_names_val: Sequence[Sequence[str]], file_names_val: Sequence[Sequence[str]],
@ -340,38 +340,43 @@ def _ssd_train_get_generators(load_data: callable,
batch_size: int, batch_size: int,
image_size: int, image_size: int,
nr_trajectories: int, nr_trajectories: int,
predictor_sizes: Sequence[Sequence[int]]) -> Tuple[Generator, int, Generator, int]: predictor_sizes: Sequence[Sequence[int]]) -> Tuple[Generator, int, Generator, Generator, int, Generator]:
if nr_trajectories == -1: if nr_trajectories == -1:
nr_trajectories = None nr_trajectories = None
train_generator, train_length = \
train_generator, train_length, train_debug_generator = \
load_data(file_names_train, instances_train, coco_path, load_data(file_names_train, instances_train, coco_path,
predictor_sizes=predictor_sizes, predictor_sizes=predictor_sizes,
batch_size=batch_size, batch_size=batch_size,
image_size=image_size, image_size=image_size,
training=True, evaluation=True, augment=False, training=True, evaluation=False, augment=False,
debug=args.debug,
nr_trajectories=nr_trajectories) nr_trajectories=nr_trajectories)
val_generator, val_length = \ val_generator, val_length, val_debug_generator = \
load_data(file_names_val, instances_val, coco_path, load_data(file_names_val, instances_val, coco_path,
predictor_sizes=predictor_sizes, predictor_sizes=predictor_sizes,
batch_size=batch_size, batch_size=batch_size,
image_size=image_size, image_size=image_size,
training=False, evaluation=False, augment=False, training=False, evaluation=False, augment=False,
debug=args.debug,
nr_trajectories=nr_trajectories) nr_trajectories=nr_trajectories)
return train_generator, train_length, val_generator, val_length return (
train_generator, train_length, train_debug_generator,
val_generator, val_length, val_debug_generator
)
def _ssd_debug_save_images(args: argparse.Namespace, save_images_on_debug: bool, def _ssd_debug_save_images(args: argparse.Namespace, save_images_on_debug: bool,
save_images: callable, get_coco_cat_maps_func: callable, save_images: callable, get_coco_cat_maps_func: callable,
summary_path: str, coco_path: str, summary_path: str, coco_path: str,
batch_size: int, image_size: int, image_size: int,
train_generator: Generator, train_length: int) -> int: train_generator: Generator) -> None:
if args.debug and save_images_on_debug: if args.debug and save_images_on_debug:
train_data = next(train_generator) train_data = next(train_generator)
train_length -= batch_size
train_images = train_data[0] train_images = train_data[0]
train_labels = train_data[1] train_labels = train_data[1]
train_labels_not_encoded = train_data[2] train_labels_not_encoded = train_data[2]
@ -383,8 +388,6 @@ def _ssd_debug_save_images(args: argparse.Namespace, save_images_on_debug: bool,
save_images(train_images, train_labels, save_images(train_images, train_labels,
summary_path, coco_path, image_size, summary_path, coco_path, image_size,
get_coco_cat_maps_func, "after-encoding") get_coco_cat_maps_func, "after-encoding")
return train_length
def _ssd_get_tensorboard_callback(args: argparse.Namespace, save_summaries_on_debug: bool, def _ssd_get_tensorboard_callback(args: argparse.Namespace, save_summaries_on_debug: bool,

View File

@ -23,7 +23,7 @@ Functions:
load_scenenet_data(...): loads the SceneNet RGB-D data into a Tensorflow data set load_scenenet_data(...): loads the SceneNet RGB-D data into a Tensorflow data set
prepare_scenenet_data(...): prepares the SceneNet RGB-D data and returns it in Python format prepare_scenenet_data(...): prepares the SceneNet RGB-D data and returns it in Python format
""" """
from typing import Callable, List, Mapping, Tuple, Optional from typing import Callable, List, Mapping, Tuple, Optional, Generator
from typing import Sequence from typing import Sequence
import numpy as np import numpy as np
@ -240,8 +240,9 @@ def load_scenenet_data(photo_paths: Sequence[Sequence[str]],
training: bool, training: bool,
evaluation: bool, evaluation: bool,
augment: bool, augment: bool,
debug: bool,
predictor_sizes: Optional[np.ndarray], predictor_sizes: Optional[np.ndarray],
nr_trajectories: Optional[int] = None) -> Tuple[callable, int]: nr_trajectories: Optional[int] = None) -> Tuple[Generator, int, Optional[Generator]]:
""" """
Loads the SceneNet RGB-D data and returns a data set. Loads the SceneNet RGB-D data and returns a data set.
@ -254,12 +255,14 @@ def load_scenenet_data(photo_paths: Sequence[Sequence[str]],
training: True if training data is desired training: True if training data is desired
evaluation: True if evaluation-ready data is desired evaluation: True if evaluation-ready data is desired
augment: True if training data should be augmented augment: True if training data should be augmented
debug: True if a more extensive generator should be added to output
predictor_sizes: sizes of the predictor layers, can be None for evaluation predictor_sizes: sizes of the predictor layers, can be None for evaluation
nr_trajectories: number of trajectories to consider nr_trajectories: number of trajectories to consider
Returns: Returns:
scenenet data set generator scenenet data set generator
length of dataset length of dataset
generator which offers processed_labels as well (only if debug is True)
""" """
trajectories = zip(photo_paths, instances) trajectories = zip(photo_paths, instances)
final_image_paths = [] final_image_paths = []
@ -320,8 +323,7 @@ def load_scenenet_data(photo_paths: Sequence[Sequence[str]],
] ]
returns = {'processed_images', 'encoded_labels'} returns = {'processed_images', 'encoded_labels'}
if training and evaluation: returns_debug = {'processed_images', 'encoded_labels', 'processed_labels'}
returns = {'processed_images', 'encoded_labels', 'processed_labels'}
if not training and evaluation: if not training and evaluation:
returns = { returns = {
@ -356,9 +358,21 @@ def load_scenenet_data(photo_paths: Sequence[Sequence[str]],
keep_images_without_gt=False keep_images_without_gt=False
) )
if debug:
debug_generator = data_generator.generate(
batch_size=batch_size,
shuffle=shuffle,
transformations=transformations,
label_encoder=label_encoder,
returns=returns_debug,
keep_images_without_gt=False
)
else:
debug_generator = None
length_dataset = data_generator.dataset_size length_dataset = data_generator.dataset_size
return generator, length_dataset return generator, length_dataset, debug_generator
def _load_images_ssd_callback(resized_shape: Sequence[int]) \ def _load_images_ssd_callback(resized_shape: Sequence[int]) \