def test_normalize(input_shape):
mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]
transfo = T.Normalize(mean, std)
input_t = tf.cast(tf.fill(input_shape, 1), dtype=tf.float32)
out = transfo(input_t)
assert tf.reduce_all(out == 1)
assert repr(transfo) == f"Normalize(mean={mean}, std={std})"
# FP16
input_t = tf.cast(tf.fill(input_shape, 1), dtype=tf.float16)
out = transfo(input_t)
assert out.dtype == tf.float16
def main(args):
print(args)
if not isinstance(args.workers, int):
args.workers = min(16, mp.cpu_count())
# AMP
if args.amp:
mixed_precision.set_global_policy("mixed_float16")
input_shape = (args.size, args.size,
3) if isinstance(args.size, int) else None
# Load docTR model
model = detection.__dict__[args.arch](
pretrained=isinstance(args.resume, str),
assume_straight_pages=not args.rotation,
input_shape=input_shape,
)
# Resume weights
if isinstance(args.resume, str):
print(f"Resuming {args.resume}")
model.load_weights(args.resume).expect_partial()
input_shape = model.cfg[
"input_shape"] if input_shape is None else input_shape
mean, std = model.cfg["mean"], model.cfg["std"]
st = time.time()
ds = datasets.__dict__[args.dataset](
train=True,
download=True,
use_polygons=args.rotation,
sample_transforms=T.Resize(input_shape[:2]),
)
# Monkeypatch
subfolder = ds.root.split("/")[-2:]
ds.root = str(Path(ds.root).parent.parent)
ds.data = [(os.path.join(*subfolder, name), target)
for name, target in ds.data]
_ds = datasets.__dict__[args.dataset](
train=False,
download=True,
use_polygons=args.rotation,
sample_transforms=T.Resize(input_shape[:2]),
)
subfolder = _ds.root.split("/")[-2:]
ds.data.extend([(os.path.join(*subfolder, name), target)
for name, target in _ds.data])
test_loader = DataLoader(
ds,
batch_size=args.batch_size,
drop_last=False,
num_workers=args.workers,
shuffle=False,
)
print(f"Test set loaded in {time.time() - st:.4}s ({len(ds)} samples in "
f"{len(test_loader)} batches)")
batch_transforms = T.Normalize(mean=mean, std=std)
# Metrics
metric = LocalizationConfusion(use_polygons=args.rotation,
mask_shape=input_shape[:2])
print("Running evaluation")
val_loss, recall, precision, mean_iou = evaluate(model, test_loader,
batch_transforms, metric)
print(
f"Validation loss: {val_loss:.6} (Recall: {recall:.2%} | Precision: {precision:.2%} | "
f"Mean IoU: {mean_iou:.2%})")
def main(args):
print(args)
if args.push_to_hub:
login_to_hub()
if not isinstance(args.workers, int):
args.workers = min(16, mp.cpu_count())
vocab = VOCABS[args.vocab]
fonts = args.font.split(",")
# AMP
if args.amp:
mixed_precision.set_global_policy("mixed_float16")
# Load val data generator
st = time.time()
val_set = CharacterGenerator(
vocab=vocab,
num_samples=args.val_samples * len(vocab),
cache_samples=True,
img_transforms=T.Compose(
[
T.Resize((args.input_size, args.input_size)),
# Ensure we have a 90% split of white-background images
T.RandomApply(T.ColorInversion(), 0.9),
]
),
font_family=fonts,
)
val_loader = DataLoader(
val_set,
batch_size=args.batch_size,
shuffle=False,
drop_last=False,
num_workers=args.workers,
collate_fn=collate_fn,
)
print(
f"Validation set loaded in {time.time() - st:.4}s ({len(val_set)} samples in "
f"{val_loader.num_batches} batches)"
)
# Load doctr model
model = classification.__dict__[args.arch](
pretrained=args.pretrained,
input_shape=(args.input_size, args.input_size, 3),
num_classes=len(vocab),
classes=list(vocab),
include_top=True,
)
# Resume weights
if isinstance(args.resume, str):
model.load_weights(args.resume)
batch_transforms = T.Compose(
[
T.Normalize(mean=(0.694, 0.695, 0.693), std=(0.299, 0.296, 0.301)),
]
)
if args.test_only:
print("Running evaluation")
val_loss, acc = evaluate(model, val_loader, batch_transforms)
print(f"Validation loss: {val_loss:.6} (Acc: {acc:.2%})")
return
st = time.time()
# Load train data generator
train_set = CharacterGenerator(
vocab=vocab,
num_samples=args.train_samples * len(vocab),
cache_samples=True,
img_transforms=T.Compose(
[
T.Resize((args.input_size, args.input_size)),
# Augmentations
T.RandomApply(T.ColorInversion(), 0.9),
T.RandomApply(T.ToGray(3), 0.1),
T.RandomJpegQuality(60),
T.RandomSaturation(0.3),
T.RandomContrast(0.3),
T.RandomBrightness(0.3),
# Blur
T.RandomApply(T.GaussianBlur(kernel_shape=(3, 3), std=(0.1, 3)), 0.3),
]
),
font_family=fonts,
)
train_loader = DataLoader(
train_set,
batch_size=args.batch_size,
shuffle=True,
drop_last=True,
num_workers=args.workers,
collate_fn=collate_fn,
)
print(
f"Train set loaded in {time.time() - st:.4}s ({len(train_set)} samples in "
f"{train_loader.num_batches} batches)"
)
if args.show_samples:
x, target = next(iter(train_loader))
plot_samples(x, list(map(vocab.__getitem__, target)))
return
# Optimizer
scheduler = tf.keras.optimizers.schedules.ExponentialDecay(
args.lr,
decay_steps=args.epochs * len(train_loader),
decay_rate=1 / (1e3), # final lr as a fraction of initial lr
staircase=False,
)
optimizer = tf.keras.optimizers.Adam(
learning_rate=scheduler,
beta_1=0.95,
beta_2=0.99,
epsilon=1e-6,
)
if args.amp:
optimizer = mixed_precision.LossScaleOptimizer(optimizer)
# LR Finder
if args.find_lr:
lrs, losses = record_lr(model, train_loader, batch_transforms, optimizer, amp=args.amp)
plot_recorder(lrs, losses)
return
# Tensorboard to monitor training
current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
exp_name = f"{args.arch}_{current_time}" if args.name is None else args.name
# W&B
if args.wb:
run = wandb.init(
name=exp_name,
project="character-classification",
config={
"learning_rate": args.lr,
"epochs": args.epochs,
"weight_decay": 0.0,
"batch_size": args.batch_size,
"architecture": args.arch,
"input_size": args.input_size,
"optimizer": "adam",
"framework": "tensorflow",
"vocab": args.vocab,
"scheduler": "exp_decay",
"pretrained": args.pretrained,
},
)
# Create loss queue
min_loss = np.inf
# Training loop
mb = master_bar(range(args.epochs))
for epoch in mb:
fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb, args.amp)
# Validation loop at the end of each epoch
val_loss, acc = evaluate(model, val_loader, batch_transforms)
if val_loss < min_loss:
print(f"Validation loss decreased {min_loss:.6} --> {val_loss:.6}: saving state...")
model.save_weights(f"./{exp_name}/weights")
min_loss = val_loss
mb.write(f"Epoch {epoch + 1}/{args.epochs} - Validation loss: {val_loss:.6} (Acc: {acc:.2%})")
# W&B
if args.wb:
wandb.log(
{
"val_loss": val_loss,
"acc": acc,
}
)
if args.wb:
run.finish()
if args.push_to_hub:
push_to_hf_hub(model, exp_name, task="classification", run_config=args)
if args.export_onnx:
print("Exporting model to ONNX...")
dummy_input = [tf.TensorSpec([None, args.input_size, args.input_size, 3], tf.float32, name="input")]
model_path, _ = export_model_to_onnx(model, exp_name, dummy_input)
print(f"Exported model saved in {model_path}")
def main(args):
print(args)
if not isinstance(args.workers, int):
args.workers = min(16, mp.cpu_count())
# AMP
if args.amp:
mixed_precision.set_global_policy("mixed_float16")
# Load doctr model
model = recognition.__dict__[args.arch](
pretrained=True if args.resume is None else False,
input_shape=(args.input_size, 4 * args.input_size, 3),
vocab=VOCABS[args.vocab],
)
# Resume weights
if isinstance(args.resume, str):
model.load_weights(args.resume)
st = time.time()
ds = datasets.__dict__[args.dataset](
train=True,
download=True,
recognition_task=True,
use_polygons=args.regular,
img_transforms=T.Resize((args.input_size, 4 * args.input_size),
preserve_aspect_ratio=True),
)
_ds = datasets.__dict__[args.dataset](
train=False,
download=True,
recognition_task=True,
use_polygons=args.regular,
img_transforms=T.Resize((args.input_size, 4 * args.input_size),
preserve_aspect_ratio=True),
)
ds.data.extend([(np_img, target) for np_img, target in _ds.data])
test_loader = DataLoader(
ds,
batch_size=args.batch_size,
drop_last=False,
num_workers=args.workers,
shuffle=False,
)
print(f"Test set loaded in {time.time() - st:.4}s ({len(ds)} samples in "
f"{len(test_loader)} batches)")
mean, std = model.cfg["mean"], model.cfg["std"]
batch_transforms = T.Normalize(mean=mean, std=std)
# Metrics
val_metric = TextMatch()
print("Running evaluation")
val_loss, exact_match, partial_match = evaluate(model, test_loader,
batch_transforms,
val_metric)
print(
f"Validation loss: {val_loss:.6} (Exact: {exact_match:.2%} | Partial: {partial_match:.2%})"
)
def main(args):
print(args)
if args.push_to_hub:
login_to_hub()
if not isinstance(args.workers, int):
args.workers = min(16, mp.cpu_count())
vocab = VOCABS[args.vocab]
fonts = args.font.split(",")
# AMP
if args.amp:
mixed_precision.set_global_policy("mixed_float16")
st = time.time()
if isinstance(args.val_path, str):
with open(os.path.join(args.val_path, "labels.json"), "rb") as f:
val_hash = hashlib.sha256(f.read()).hexdigest()
# Load val data generator
val_set = RecognitionDataset(
img_folder=os.path.join(args.val_path, "images"),
labels_path=os.path.join(args.val_path, "labels.json"),
img_transforms=T.Resize((args.input_size, 4 * args.input_size),
preserve_aspect_ratio=True),
)
else:
val_hash = None
# Load synthetic data generator
val_set = WordGenerator(
vocab=vocab,
min_chars=args.min_chars,
max_chars=args.max_chars,
num_samples=args.val_samples * len(vocab),
font_family=fonts,
img_transforms=T.Compose([
T.Resize((args.input_size, 4 * args.input_size),
preserve_aspect_ratio=True),
# Ensure we have a 90% split of white-background images
T.RandomApply(T.ColorInversion(), 0.9),
]),
)
val_loader = DataLoader(
val_set,
batch_size=args.batch_size,
shuffle=False,
drop_last=False,
num_workers=args.workers,
)
print(
f"Validation set loaded in {time.time() - st:.4}s ({len(val_set)} samples in "
f"{val_loader.num_batches} batches)")
# Load doctr model
model = recognition.__dict__[args.arch](
pretrained=args.pretrained,
input_shape=(args.input_size, 4 * args.input_size, 3),
vocab=vocab,
)
# Resume weights
if isinstance(args.resume, str):
model.load_weights(args.resume)
# Metrics
val_metric = TextMatch()
batch_transforms = T.Compose([
T.Normalize(mean=(0.694, 0.695, 0.693), std=(0.299, 0.296, 0.301)),
])
if args.test_only:
print("Running evaluation")
val_loss, exact_match, partial_match = evaluate(
model, val_loader, batch_transforms, val_metric)
print(
f"Validation loss: {val_loss:.6} (Exact: {exact_match:.2%} | Partial: {partial_match:.2%})"
)
return
st = time.time()
if isinstance(args.train_path, str):
# Load train data generator
base_path = Path(args.train_path)
parts = ([base_path]
if base_path.joinpath("labels.json").is_file() else
[base_path.joinpath(sub) for sub in os.listdir(base_path)])
with open(parts[0].joinpath("labels.json"), "rb") as f:
train_hash = hashlib.sha256(f.read()).hexdigest()
train_set = RecognitionDataset(
parts[0].joinpath("images"),
parts[0].joinpath("labels.json"),
img_transforms=T.Compose([
T.RandomApply(T.ColorInversion(), 0.1),
T.Resize((args.input_size, 4 * args.input_size),
preserve_aspect_ratio=True),
# Augmentations
T.RandomJpegQuality(60),
T.RandomSaturation(0.3),
T.RandomContrast(0.3),
T.RandomBrightness(0.3),
]),
)
if len(parts) > 1:
for subfolder in parts[1:]:
train_set.merge_dataset(
RecognitionDataset(subfolder.joinpath("images"),
subfolder.joinpath("labels.json")))
else:
train_hash = None
# Load synthetic data generator
train_set = WordGenerator(
vocab=vocab,
min_chars=args.min_chars,
max_chars=args.max_chars,
num_samples=args.train_samples * len(vocab),
font_family=fonts,
img_transforms=T.Compose([
T.Resize((args.input_size, 4 * args.input_size),
preserve_aspect_ratio=True),
# Ensure we have a 90% split of white-background images
T.RandomApply(T.ColorInversion(), 0.9),
T.RandomJpegQuality(60),
T.RandomSaturation(0.3),
T.RandomContrast(0.3),
T.RandomBrightness(0.3),
]),
)
train_loader = DataLoader(
train_set,
batch_size=args.batch_size,
shuffle=True,
drop_last=True,
num_workers=args.workers,
)
print(
f"Train set loaded in {time.time() - st:.4}s ({len(train_set)} samples in "
f"{train_loader.num_batches} batches)")
if args.show_samples:
x, target = next(iter(train_loader))
plot_samples(x, target)
return
# Optimizer
scheduler = tf.keras.optimizers.schedules.ExponentialDecay(
args.lr,
decay_steps=args.epochs * len(train_loader),
decay_rate=1 / (25e4), # final lr as a fraction of initial lr
staircase=False,
)
optimizer = tf.keras.optimizers.Adam(learning_rate=scheduler,
beta_1=0.95,
beta_2=0.99,
epsilon=1e-6,
clipnorm=5)
if args.amp:
optimizer = mixed_precision.LossScaleOptimizer(optimizer)
# LR Finder
if args.find_lr:
lrs, losses = record_lr(model,
train_loader,
batch_transforms,
optimizer,
amp=args.amp)
plot_recorder(lrs, losses)
return
# Tensorboard to monitor training
current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
exp_name = f"{args.arch}_{current_time}" if args.name is None else args.name
# W&B
if args.wb:
run = wandb.init(
name=exp_name,
project="text-recognition",
config={
"learning_rate": args.lr,
"epochs": args.epochs,
"weight_decay": 0.0,
"batch_size": args.batch_size,
"architecture": args.arch,
"input_size": args.input_size,
"optimizer": "adam",
"framework": "tensorflow",
"scheduler": "exp_decay",
"vocab": args.vocab,
"train_hash": train_hash,
"val_hash": val_hash,
"pretrained": args.pretrained,
},
)
min_loss = np.inf
# Training loop
mb = master_bar(range(args.epochs))
for epoch in mb:
fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb,
args.amp)
# Validation loop at the end of each epoch
val_loss, exact_match, partial_match = evaluate(
model, val_loader, batch_transforms, val_metric)
if val_loss < min_loss:
print(
f"Validation loss decreased {min_loss:.6} --> {val_loss:.6}: saving state..."
)
model.save_weights(f"./{exp_name}/weights")
min_loss = val_loss
mb.write(
f"Epoch {epoch + 1}/{args.epochs} - Validation loss: {val_loss:.6} "
f"(Exact: {exact_match:.2%} | Partial: {partial_match:.2%})")
# W&B
if args.wb:
wandb.log({
"val_loss": val_loss,
"exact_match": exact_match,
"partial_match": partial_match,
})
if args.wb:
run.finish()
if args.push_to_hub:
push_to_hf_hub(model, exp_name, task="recognition", run_config=args)
def main(args):
print(args)
if args.push_to_hub:
login_to_hub()
if not isinstance(args.workers, int):
args.workers = min(16, mp.cpu_count())
# AMP
if args.amp:
mixed_precision.set_global_policy("mixed_float16")
st = time.time()
val_set = DetectionDataset(
img_folder=os.path.join(args.val_path, "images"),
label_path=os.path.join(args.val_path, "labels.json"),
sample_transforms=T.SampleCompose(([
T.Resize((args.input_size, args.input_size),
preserve_aspect_ratio=True,
symmetric_pad=True)
] if not args.rotation or args.eval_straight else []) + ([
T.Resize(args.input_size, preserve_aspect_ratio=True
), # This does not pad
T.RandomRotate(90, expand=True),
T.Resize((args.input_size, args.input_size),
preserve_aspect_ratio=True,
symmetric_pad=True),
] if args.rotation and not args.eval_straight else [])),
use_polygons=args.rotation and not args.eval_straight,
)
val_loader = DataLoader(
val_set,
batch_size=args.batch_size,
shuffle=False,
drop_last=False,
num_workers=args.workers,
)
print(
f"Validation set loaded in {time.time() - st:.4}s ({len(val_set)} samples in "
f"{val_loader.num_batches} batches)")
with open(os.path.join(args.val_path, "labels.json"), "rb") as f:
val_hash = hashlib.sha256(f.read()).hexdigest()
batch_transforms = T.Compose([
T.Normalize(mean=(0.798, 0.785, 0.772), std=(0.264, 0.2749, 0.287)),
])
# Load doctr model
model = detection.__dict__[args.arch](
pretrained=args.pretrained,
input_shape=(args.input_size, args.input_size, 3),
assume_straight_pages=not args.rotation,
)
# Resume weights
if isinstance(args.resume, str):
model.load_weights(args.resume)
# Metrics
val_metric = LocalizationConfusion(use_polygons=args.rotation
and not args.eval_straight,
mask_shape=(args.input_size,
args.input_size))
if args.test_only:
print("Running evaluation")
val_loss, recall, precision, mean_iou = evaluate(
model, val_loader, batch_transforms, val_metric)
print(
f"Validation loss: {val_loss:.6} (Recall: {recall:.2%} | Precision: {precision:.2%} | "
f"Mean IoU: {mean_iou:.2%})")
return
st = time.time()
# Load both train and val data generators
train_set = DetectionDataset(
img_folder=os.path.join(args.train_path, "images"),
label_path=os.path.join(args.train_path, "labels.json"),
img_transforms=T.Compose([
# Augmentations
T.RandomApply(T.ColorInversion(), 0.1),
T.RandomJpegQuality(60),
T.RandomSaturation(0.3),
T.RandomContrast(0.3),
T.RandomBrightness(0.3),
]),
sample_transforms=T.SampleCompose(([
T.Resize((args.input_size, args.input_size),
preserve_aspect_ratio=True,
symmetric_pad=True)
] if not args.rotation else []) + ([
T.Resize(args.input_size, preserve_aspect_ratio=True
), # This does not pad
T.RandomRotate(90, expand=True),
T.Resize((args.input_size, args.input_size),
preserve_aspect_ratio=True,
symmetric_pad=True),
] if args.rotation else [])),
use_polygons=args.rotation,
)
train_loader = DataLoader(
train_set,
batch_size=args.batch_size,
shuffle=True,
drop_last=True,
num_workers=args.workers,
)
print(
f"Train set loaded in {time.time() - st:.4}s ({len(train_set)} samples in "
f"{train_loader.num_batches} batches)")
with open(os.path.join(args.train_path, "labels.json"), "rb") as f:
train_hash = hashlib.sha256(f.read()).hexdigest()
if args.show_samples:
x, target = next(iter(train_loader))
plot_samples(x, target)
return
# Optimizer
scheduler = tf.keras.optimizers.schedules.ExponentialDecay(
args.lr,
decay_steps=args.epochs * len(train_loader),
decay_rate=1 / (25e4), # final lr as a fraction of initial lr
staircase=False,
)
optimizer = tf.keras.optimizers.Adam(learning_rate=scheduler,
beta_1=0.95,
beta_2=0.99,
epsilon=1e-6,
clipnorm=5)
if args.amp:
optimizer = mixed_precision.LossScaleOptimizer(optimizer)
# LR Finder
if args.find_lr:
lrs, losses = record_lr(model,
train_loader,
batch_transforms,
optimizer,
amp=args.amp)
plot_recorder(lrs, losses)
return
# Tensorboard to monitor training
current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
exp_name = f"{args.arch}_{current_time}" if args.name is None else args.name
# W&B
if args.wb:
run = wandb.init(
name=exp_name,
project="text-detection",
config={
"learning_rate": args.lr,
"epochs": args.epochs,
"weight_decay": 0.0,
"batch_size": args.batch_size,
"architecture": args.arch,
"input_size": args.input_size,
"optimizer": "adam",
"framework": "tensorflow",
"scheduler": "exp_decay",
"train_hash": train_hash,
"val_hash": val_hash,
"pretrained": args.pretrained,
"rotation": args.rotation,
},
)
if args.freeze_backbone:
for layer in model.feat_extractor.layers:
layer.trainable = False
min_loss = np.inf
# Training loop
mb = master_bar(range(args.epochs))
for epoch in mb:
fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb,
args.amp)
# Validation loop at the end of each epoch
val_loss, recall, precision, mean_iou = evaluate(
model, val_loader, batch_transforms, val_metric)
if val_loss < min_loss:
print(
f"Validation loss decreased {min_loss:.6} --> {val_loss:.6}: saving state..."
)
model.save_weights(f"./{exp_name}/weights")
min_loss = val_loss
log_msg = f"Epoch {epoch + 1}/{args.epochs} - Validation loss: {val_loss:.6} "
if any(val is None for val in (recall, precision, mean_iou)):
log_msg += "(Undefined metric value, caused by empty GTs or predictions)"
else:
log_msg += f"(Recall: {recall:.2%} | Precision: {precision:.2%} | Mean IoU: {mean_iou:.2%})"
mb.write(log_msg)
# W&B
if args.wb:
wandb.log({
"val_loss": val_loss,
"recall": recall,
"precision": precision,
"mean_iou": mean_iou,
})
if args.wb:
run.finish()
if args.push_to_hub:
push_to_hf_hub(model, exp_name, task="detection", run_config=args)