parser = argparse.ArgumentParser()
parser.add_argument("--epoch",
help="Epoch",
type=int,
required=False,
default=25)
parser.add_argument("--image",
help="Test image number",
type=int,
required=False,
default=0)
args = parser.parse_args()
# Set up the model and load the weights at the chosen epoch
transcriber = transcriberModel()
weights_dir = ("%s/Robot_Rainfall_Rescue/models/ATB2_retuned/Epoch_%04d") % (
os.getenv("SCRATCH"),
args.epoch - 1,
)
load_status = transcriber.load_weights("%s/ckpt" % weights_dir)
# Check the load worked
load_status.assert_existing_objects_matched()
# Get test case number args.image
testImage = getImageDataset(subdir="unperturbed",
purpose="test",
nImages=args.image + 1)
testImage = testImage.batch(1)
originalImage = next(itertools.islice(testImage, args.image, args.image + 1))
testNumbers = getNumbersDataset(subdir="unperturbed",
linewidth=0.0,
frameon=False,
subplotpars=None,
tight_layout=None,
)
matplotlib.rcParams.update({"font.size": 22})
canvas = FigureCanvas(fig)
# Paint the background white - why is this needed?
ax_full = fig.add_axes([0, 0, 1, 1])
ax_full.set_xlim([0, 1])
ax_full.set_ylim([0, 1])
ax_full.add_patch(
matplotlib.patches.Rectangle((0, 0), 1, 1, fill=True, facecolor="white"))
# Set up the model and load the weights at the chosen epoch
transcriber1 = transcriberModel()
weights_dir = ("%s/Robot_Rainfall_Rescue/models/ATB2_retuned/" +
"Epoch_%04d") % (
os.getenv("SCRATCH"),
int(args.epoch) - 1,
)
if args.epoch >= 1:
load_status = transcriber1.load_weights("%s/ckpt" % weights_dir)
load_status.assert_existing_objects_matched()
transcriber2 = transcriberModel()
weights_dir = ("%s/Robot_Rainfall_Rescue/models/ATB2_retuned/" +
"Epoch_%04d") % (
os.getenv("SCRATCH"),
int(args.epoch),
)