ship-vs-iceberg/plot_augmented.py
2021-10-16 08:49:48 +02:00

134 lines
4.1 KiB
Python
Executable File

#!/usr/bin/env python3
"""Plot image produced by data generators in icenet/utils.py"""
import os
import argparse
import json
import random
import numpy as np
import matplotlib.pyplot as plt
from icenet import util
from plot import plot_sample
from itertools import islice
def plot_base_cnn_samples(outdir, samples, interactive=False):
m = len(samples)
minmax = util.get_minmax(samples)
imgs, _ = list(islice(util.base_cnn_generator(
samples, minmax, m, shuffle=False, verbose=True
), 1))[0]
for i in range(m):
print("Plotting sample %d/%d" % (i+1, m))
sample = samples[i]
img = imgs[i, :, :, :]
fig, (ax1, ax2) = plt.subplots(1, 2, sharex='col', sharey='row')
fig.suptitle('\n'.join([
"Id: %s (rotated and cropped)" % sample.get('id'),
"Iceberg? %s" % sample.get('is_iceberg', '-'),
"Incident angle: %s" % repr(sample.get('inc_angle'))])
)
ax1.imshow(img[:, :, 0], vmin=0.2, vmax=0.8, aspect='auto')
ax1.set_title("HH")
ax2.imshow(img[:, :, 1], vmin=0.2, vmax=0.8, aspect='auto')
ax2.set_title("HV")
#fig.tight_layout()
fig.subplots_adjust(top=0.80)
plot_sample(outdir, sample, minmax, interactive=interactive)
if not interactive:
fig.savefig(os.path.join(outdir, "%s_base.png" % sample['id']))
plt.close('all')
def plot_icenet_samples(outdir, samples, interactive=False):
m = len(samples)
minmax = util.get_minmax(samples)
sections, y = list(islice(util.icenet_generator(
samples, minmax, m, augment=False, verbose=True
), 1))[0]
for i in range(m):
print("Plotting sample %d/%d" % (i+1, m))
sample = samples[i]
for b in range(2):
fig, axes = plt.subplots(3, 3, sharex='col', sharey='row')
fig.suptitle('\n'.join([
"Id: %s (rotated and cropped, band %d)" % (sample.get('id'), b+1),
"Iceberg? %s" % sample.get('is_iceberg', '-'),
"Incident angle: %s" % repr(sample.get('inc_angle'))])
)
for ridx in range(3):
for cidx in range(3):
axes[ridx, cidx].imshow(
sections[ridx*3+cidx][i, :, :, b],
vmin=0.2,
vmax=0.8,
aspect='auto'
)
fig.subplots_adjust(top=0.80)
if not interactive:
fig.savefig(os.path.join(outdir, "%s_icenet_b%d.png" % (sample['id'], b)))
plot_sample(outdir, sample, minmax, interactive=interactive)
plt.close('all')
if __name__ == '__main__':
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
'-o',
required=False,
help="Output directory for samples."
)
parser.add_argument(
'-g',
required=True,
choices=['base', 'icenet'],
help="Generator to use (base|icenet)."
)
parser.add_argument(
'-n',
type=int,
required=False,
help="Number of samples to plot (randomized)"
)
parser.add_argument(
'-i',
required=False,
action='store_true',
help="Show plots (and don't store)"
)
parser.add_argument(
'samples_file',
help="JSON file with samples"
)
args = parser.parse_args()
outdir = args.o or '.'
if not os.path.isdir(outdir):
raise Exception("Output directory does not exist")
print("Loading samples ...")
with open(args.samples_file) as f:
samples = json.load(f)
print("%d samples in set" % len(samples))
if args.n:
print("Pick %d random samples ..." % args.n)
random.shuffle(samples)
samples = samples[0:args.n]
if args.g == 'base':
plot_base_cnn_samples(outdir, samples, interactive=args.i or False)
else:
plot_icenet_samples(outdir, samples, interactive=args.i or False)
print("Done.")