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