def __init__(self, settings: TrainSettings): self.settings = settings tf.keras.backend.clear_session() if settings.gpu_allow_growth: from ocr4all_pixel_classifier.lib.network import tf_backend_allow_growth tf_backend_allow_growth() from ocr4all_pixel_classifier.lib.network import Network self.train_net = Network("train", settings.n_classes, settings.architecture, l_rate=settings.l_rate, foreground_masks=settings.foreground_masks, model=settings.load, continue_training=settings.continue_training, input_image_dimension=settings.image_dimension, optimizer=settings.optimizer, optimizer_norm_clipping=settings.optimizer_norm_clipping, optimizer_norm_clip_value=settings.optimizer_norm_clip_value, optimizer_clipping=settings.optimizer_clipping, optimizer_clip_value=settings.optimizer_clip_value, loss_func=settings.loss ) if len(settings.train_data) == 0 and settings.n_epoch > 0: raise Exception("No training files specified. Maybe set n_iter=0") if settings.compute_baseline: def compute_label_percentage(label): return np.sum([np.sum(d.mask == label) for d in settings.train_data.data]) \ / np.sum([d.mask.shape[0] * d.mask.shape[1] for d in settings.train_data.data]) logging.info("Computing label percentage for {} files.".format(len(settings.train_data.data))) label_percentage = [compute_label_percentage(l) for l in range(settings.n_classes)] logging.info("Label percentage: {}".format(list(zip(range(settings.n_classes), label_percentage)))) logging.info("Baseline: {}".format(max(label_percentage)))
class Trainer: def __init__(self, settings: TrainSettings): self.settings = settings tf.keras.backend.clear_session() if settings.gpu_allow_growth: from ocr4all_pixel_classifier.lib.network import tf_backend_allow_growth tf_backend_allow_growth() from ocr4all_pixel_classifier.lib.network import Network self.train_net = Network("train", settings.n_classes, settings.architecture, l_rate=settings.l_rate, foreground_masks=settings.foreground_masks, model=settings.load, continue_training=settings.continue_training, input_image_dimension=settings.image_dimension, optimizer=settings.optimizer, optimizer_norm_clipping=settings.optimizer_norm_clipping, optimizer_norm_clip_value=settings.optimizer_norm_clip_value, optimizer_clipping=settings.optimizer_clipping, optimizer_clip_value=settings.optimizer_clip_value, loss_func=settings.loss ) if len(settings.train_data) == 0 and settings.n_epoch > 0: raise Exception("No training files specified. Maybe set n_iter=0") if settings.compute_baseline: def compute_label_percentage(label): return np.sum([np.sum(d.mask == label) for d in settings.train_data.data]) \ / np.sum([d.mask.shape[0] * d.mask.shape[1] for d in settings.train_data.data]) logging.info("Computing label percentage for {} files.".format(len(settings.train_data.data))) label_percentage = [compute_label_percentage(l) for l in range(settings.n_classes)] logging.info("Label percentage: {}".format(list(zip(range(settings.n_classes), label_percentage)))) logging.info("Baseline: {}".format(max(label_percentage))) def train(self, callback: Optional[TrainProgressCallback] = None) -> None: if callback: callback.init(self.settings.n_epoch * len(self.settings.train_data.data), self.settings.early_stopping_max_performance_drops) self.train_net.train_dataset(setting=self.settings, callback=callback) def eval(self) -> None: if self.settings.evaluation_data is None: logger.info('Evaluation Dataset in Trainsetting not set! ') return if len(self.settings.evaluation_data) > 0: self.train_net.evaluate_dataset(self.settings.evaluation_data) else: logger.info('Empty Dataset. Skipping Evaluation')
class Predictor: def __init__(self, settings: PredictSettings, network: Network = None): self.settings = settings self.network = network if settings.gpu_allow_growth: tf_backend_allow_growth() if not network: self.network = Network("Predict", n_classes=settings.n_classes, model_constructor=settings.architecture, model=os.path.abspath( self.settings.network)) if settings.output: output_dir = settings.output os.makedirs(os.path.join(output_dir, "overlay"), exist_ok=True) os.makedirs(os.path.join(output_dir, "color"), exist_ok=True) os.makedirs(os.path.join(output_dir, "inverted"), exist_ok=True) def predict(self, dataset: Dataset) -> Generator[Prediction, None, None]: for data in dataset.data: prediction = self.predict_single(data) yield prediction def predict_single(self, data: SingleData) -> Prediction: logit, prob, pred = self.network.predict_single_data(data) if self.settings.high_res_output: data, pred = scale_to_original_shape(data, pred) if self.settings.post_process: for processor in self.settings.post_process: pred = processor(pred, data) return Prediction(pred, prob, data) def predict_masks(self, data: SingleData) -> Masks: logit, prob, pred = self.network.predict_single_data(data) if self.settings.high_res_output: data, pred = scale_to_original_shape(data, pred) if self.settings.post_process: for processor in self.settings.post_process: pred = processor(pred, data) return generate_output_masks(data, pred, self.settings.color_map)
def __init__(self, settings: PredictSettings, network: Network = None): self.settings = settings self.network = network if settings.gpu_allow_growth: tf_backend_allow_growth() if not network: self.network = Network("Predict", n_classes=settings.n_classes, model=os.path.abspath( self.settings.network)) if settings.output: output_dir = settings.output os.makedirs(os.path.join(output_dir, "overlay"), exist_ok=True) os.makedirs(os.path.join(output_dir, "color"), exist_ok=True) os.makedirs(os.path.join(output_dir, "inverted"), exist_ok=True)
class Predictor: def __init__(self, settings: PredictSettings, network: Network = None): self.settings = settings self.network = network if settings.gpu_allow_growth: tf_backend_allow_growth() if not network: self.network = Network("Predict", n_classes=settings.n_classes, model=os.path.abspath( self.settings.network)) if settings.output: output_dir = settings.output os.makedirs(os.path.join(output_dir, "overlay"), exist_ok=True) os.makedirs(os.path.join(output_dir, "color"), exist_ok=True) os.makedirs(os.path.join(output_dir, "inverted"), exist_ok=True) def predict(self, dataset: Dataset) -> Generator[Prediction, None, None]: for data in dataset.data: logit, prob, pred = self.network.predict_single_data(data) if self.settings.post_process: for processor in self.settings.post_process: pred = processor(pred, data) self.output_data(pred, data) yield Prediction(pred, prob, data) def predict_masks(self, data: SingleData) -> Masks: logit, prob, pred = self.network.predict_single_data(data) if self.settings.post_process: for processor in self.settings.post_process: pred = processor(pred, data) color_mask, inverted_overlay_mask, overlay_mask = self.generate_output_masks( data, pred) return Masks(color=color_mask, overlay=overlay_mask, inverted_overlay=inverted_overlay_mask) def output_data(self, pred, data: SingleData): if len(pred.shape) == 3: assert (pred.shape[0] == 1) pred = pred[0] if self.settings.output: if data.output_path: filename = data.output_path dir = os.path.dirname(filename) if os.path.isabs(dir): os.makedirs(dir, exist_ok=True) elif dir: for category in ["color", "overlay", "inverted"]: os.makedirs(os.path.join(self.settings.output, category, dir), exist_ok=True) else: filename = os.path.basename(data.image_path) color_mask, inverted_overlay_mask, overlay_mask = self.generate_output_masks( data, pred) img_io.imsave( os.path.join(self.settings.output, "color", filename), (color_mask * 255).astype(np.uint8)) img_io.imsave( os.path.join(self.settings.output, "overlay", filename), (overlay_mask * 255).astype(np.uint8)) img_io.imsave( os.path.join(self.settings.output, "inverted", filename), (inverted_overlay_mask * 255).astype(np.uint8)) def generate_output_masks(self, data, pred): from ocr4all_pixel_classifier.lib.dataset import label_to_colors color_mask = label_to_colors(pred, colormap=self.settings.color_map) foreground = np.stack([(1 - data.binary)] * 3, axis=-1) inv_binary = data.binary if self.settings.high_res_output: color_mask = resize(color_mask, data.original_shape, order=0) foreground = resize(foreground, data.original_shape) / 255 inv_binary = resize(inv_binary, data.original_shape, order=0) inv_binary = np.stack([inv_binary] * 3, axis=-1) overlay_mask = color_mask.copy() overlay_mask[foreground == 0] = 0 inverted_overlay_mask = color_mask.copy() inverted_overlay_mask[inv_binary == 0] = 0 return color_mask, inverted_overlay_mask, overlay_mask