def main(args, _run, _log):
log_config(_run, _log)
# load model and data:
log.info('Loading network and dataset')
net, datloader, img_inds = load_net_and_data(args)
if args['classindex'] is not None:
log.info('Specified class: {} // Number of images containing this class: {}'.format(args['classindex'],
len(img_inds)))
# predict all images and save outputs together with filename and annotation to hdf5 file:
log.info('Predicting images...')
n = len(datloader)
log.debug('Total number of batches to process: {}'.format(n))
for i, (img, lbl, image_path) in enumerate(datloader):
out = pred(net, img)
save(out, lbl, image_path, img_inds[i], input_dir=args['input_dir'])
if ((i + 1) % 1) == 0:
log.info('\t\t Image {}/{}'.format(i + 1, n))
def main(args, _run, _log):
log_config(_run, _log)
if args['load_file'] is None:
_log.info('Loading data...')
_log.info('Cityscapes...')
# load cityscapes train data for normalization of out of domain data
_, _, _, _, xa_mean, xa_std, classes_mean, classes_std, *_ = load_data(
'cityscapes')
_log.info('{}...'.format(args['dataset']))
xa, *_, start, _ = load_data(args['dataset'],
xa_mean=xa_mean,
xa_std=xa_std,
classes_mean=classes_mean,
classes_std=classes_std)
# predict iou using MetaSeg metrics
iou_pred = meta_nn_predict(args['meta_nn_path'], xa, gpu=args['gpu'])
# get all available input file IDs
inds = get_indices(
join(CONFIG.metaseg_io_path, 'metrics', 'deeplabv3plus',
args['dataset']))
# construct thresholds and dictionary for saving
thresholds = np.linspace(0, 1, args['steps'])
confusion_matrices = dict(pos={
t: np.zeros((num_categories, num_categories))
for t in thresholds
},
neg={
t: np.zeros(
(num_categories, num_categories))
for t in thresholds
})
with open('/data/poberdie/metaseg/confusion_matrices.p', 'rb') as f:
confusion_matrices['pos'] = pkl.load(f)
_log.info('Calculating IoUs...')
inputs = [(ind, iou_pred[start[i]:start[i + 1]], thresholds)
for i, ind in enumerate(inds)]
with Pool(args['n_workers']) as p:
res = list(
tqdm.tqdm(p.imap(pool_wrapper, inputs), total=len(inputs)))
for r in res:
# for t, v in r[0].items():
# confusion_matrices['pos'][t] += v
for t, v in r[1].items():
confusion_matrices['neg'][t] += v
with open(
join(args['save_dir'],
'confusion_matrices_{}.p'.format(args['dataset'])),
'wb') as f:
pkl.dump(confusion_matrices, f)
else:
thresholds = np.linspace(0, 1, args['steps'])
with open(args['load_file'], 'rb') as f:
confusion_matrices = pkl.load(f)
_log.info('Start plotting')
ious_pos = []
for t, v in confusion_matrices['pos'].items():
tp = np.diag(v)[1:]
fn = v.sum(0)[1:] - tp
fp = v.sum(1)[1:] - tp
ious_pos.append((t, (tp / (tp + fp + fn + 1e-6)).mean() * 100))
ious_pos.sort(key=lambda x: x[0])
ious_pos = [i[1] for i in ious_pos if i[0] < args['max_t']]
ious_neg = []
for t, v in confusion_matrices['neg'].items():
tp = np.diag(v)[1:]
fn = v.sum(0)[1:] - tp
fp = v.sum(1)[1:] - tp
ious_neg.append((t, (tp / (tp + fp + fn + 1e-6)).mean() * 100))
ious_neg.sort(key=lambda x: x[0])
ious_neg = [i[1] for i in ious_neg if i[0] < args['max_t']]
colmap = plt.get_cmap('tab20c')
fig = plt.figure('IoU under different thresholds', dpi=args['dpi'])
ax = fig.add_subplot(111)
ax.set_axisbelow(True)
ax.grid(linestyle='--')
ax.set_xlabel('Threshold t on predicted IoU')
ax.set_ylabel('mIoU in % for segments below t', color=colmap(4))
ax.tick_params(axis='y', labelcolor=colmap(4))
ax2 = ax.twinx()
ax2.set_ylabel('mIoU in % for segments above t', color=colmap(8))
ax2.tick_params(axis='y', labelcolor=colmap(8))
ax.plot(thresholds[thresholds < args['max_t']],
np.array(ious_neg),
color=colmap(4),
linestyle='-')
ax2.plot(thresholds[thresholds < args['max_t']],
np.array(ious_pos),
color=colmap(8),
linestyle='-')
plt.savefig(join(args['plot_dir'],
'iou_plot_{}.png'.format(args['dataset'])),
dpi=args['dpi'],
bbox_inches='tight')
_log.info('Saved plot to \'{}\''.format(
join(args['plot_dir'], 'iou_plot_{}.png'.format(args['dataset']))))
def main(args, _run, _log):
log_config(_run, _log)
if args["load_file"] is None:
_log.info("Loading data...")
_log.info("Cityscapes...")
# load cityscapes train data for normalization of out of domain data
_, _, _, _, xa_mean, xa_std, classes_mean, classes_std, *_ = load_data(
"cityscapes")
_log.info("{}...".format(args["dataset"]))
xa, *_, start, _ = load_data(
args["dataset"],
xa_mean=xa_mean,
xa_std=xa_std,
classes_mean=classes_mean,
classes_std=classes_std,
)
# predict iou using MetaSeg metrics
iou_pred = meta_nn_predict(args["meta_nn_path"], xa, gpu=args["gpu"])
# get all available input file IDs
inds = get_indices(
join(CONFIG.metaseg_io_path, "metrics", "deeplabv3plus",
args["dataset"]))
# construct thresholds and dictionary for saving
thresholds = np.linspace(0, 1, args["steps"])
confusion_matrices = dict(
pos={
t: np.zeros((num_categories, num_categories))
for t in thresholds
},
neg={
t: np.zeros((num_categories, num_categories))
for t in thresholds
},
)
with open("/data/poberdie/metaseg/confusion_matrices.p", "rb") as f:
confusion_matrices["pos"] = pkl.load(f)
_log.info("Calculating IoUs...")
inputs = [(ind, iou_pred[start[i]:start[i + 1]], thresholds)
for i, ind in enumerate(inds)]
with Pool(args["n_workers"]) as p:
res = list(
tqdm.tqdm(p.imap(pool_wrapper, inputs), total=len(inputs)))
for r in res:
# for t, v in r[0].items():
# confusion_matrices['pos'][t] += v
for t, v in r[1].items():
confusion_matrices["neg"][t] += v
with open(
join(args["save_dir"],
"confusion_matrices_{}.p".format(args["dataset"])),
"wb",
) as f:
pkl.dump(confusion_matrices, f)
else:
thresholds = np.linspace(0, 1, args["steps"])
with open(args["load_file"], "rb") as f:
confusion_matrices = pkl.load(f)
_log.info("Start plotting")
ious_pos = []
for t, v in confusion_matrices["pos"].items():
tp = np.diag(v)[1:]
fn = v.sum(0)[1:] - tp
fp = v.sum(1)[1:] - tp
ious_pos.append((t, (tp / (tp + fp + fn + 1e-6)).mean() * 100))
ious_pos.sort(key=lambda x: x[0])
ious_pos = [i[1] for i in ious_pos if i[0] < args["max_t"]]
ious_neg = []
for t, v in confusion_matrices["neg"].items():
tp = np.diag(v)[1:]
fn = v.sum(0)[1:] - tp
fp = v.sum(1)[1:] - tp
ious_neg.append((t, (tp / (tp + fp + fn + 1e-6)).mean() * 100))
ious_neg.sort(key=lambda x: x[0])
ious_neg = [i[1] for i in ious_neg if i[0] < args["max_t"]]
colmap = plt.get_cmap("tab20c")
fig = plt.figure("IoU under different thresholds", dpi=args["dpi"])
ax = fig.add_subplot(111)
ax.set_axisbelow(True)
ax.grid(linestyle="--")
ax.set_xlabel("Threshold t on predicted IoU")
ax.set_ylabel("mIoU in % for segments below t", color=colmap(4))
ax.tick_params(axis="y", labelcolor=colmap(4))
ax2 = ax.twinx()
ax2.set_ylabel("mIoU in % for segments above t", color=colmap(8))
ax2.tick_params(axis="y", labelcolor=colmap(8))
ax.plot(
thresholds[thresholds < args["max_t"]],
np.array(ious_neg),
color=colmap(4),
linestyle="-",
)
ax2.plot(
thresholds[thresholds < args["max_t"]],
np.array(ious_pos),
color=colmap(8),
linestyle="-",
)
plt.savefig(
join(args["plot_dir"], "iou_plot_{}.png".format(args["dataset"])),
dpi=args["dpi"],
bbox_inches="tight",
)
_log.info("Saved plot to '{}'".format(
join(args["plot_dir"], "iou_plot_{}.png".format(args["dataset"]))))
def main(args, mainplot, tsne, _run, _log):
log_config(_run, _log)
Discovery(**args, main_plot_args=mainplot, tsne_args=tsne)
def main(args, tsne, _run, _log):
log_config(_run, _log)
with open(args["embeddings_file"], "rb") as f:
data = pkl.load(f)
gt = np.array(data["gt"]).squeeze()
_log.debug("Number of segments: {}".format(gt.shape[0]))
gt = np.vectorize(id_to_trainid.get)(gt)
if (data["nn_embeddings"].shape[1] != args["embedding_size"]
if "nn_embeddings" in data.keys() else
True) or args["overwrite_embeddings"]:
embeddings = np.stack(data["embeddings"])
# _log.info('Standardizing embeddings...')
# embeddings = (embeddings - embeddings.mean()) / embeddings.std()
if (args["embedding_size"] < embeddings.shape[1]
if args["embedding_size"] is not None else False):
_log.info("Computing embeddings for nearest neighbor search...")
if args["method"] == "TSNE":
_log.info("Using t-SNE with method '{}' and dimensionality {}".
format(
"barnes_hut"
if args["embedding_size"] < 4 else "exact",
args["embedding_size"],
))
embeddings = PCA(n_components=50 if args["embedding_size"] < 50
else 100).fit_transform(embeddings)
embeddings = TSNE(n_components=args["embedding_size"],
n_jobs=args["n_jobs"],
method="barnes_hut"
if args["embedding_size"] < 4 else "exact",
**tsne).fit_transform(embeddings)
elif args["method"] == "Isomap":
_log.info("Using Isomap method.")
embeddings = PCA(n_components=50 if args["embedding_size"] < 50
else 100).fit_transform(embeddings)
embeddings = Isomap(
n_components=args["embedding_size"],
n_jobs=args["n_jobs"],
).fit_transform(embeddings)
elif args["method"] == "PCA":
_log.info("Using PCA method.")
embeddings = PCA(n_components=args["embedding_size"]
).fit_transform(embeddings)
data["nn_embeddings"] = embeddings
_log.debug("Saving computed manifold to embeddings file.")
with open(args["embeddings_file"], "wb") as f:
pkl.dump(data, f)
else:
_log.info("Leaving data as it is.")
else:
embeddings = data["nn_embeddings"]
_log.info(("Using precomputed embeddings "
"({} dimensions) for nearest neighbor search...".format(
embeddings.shape[1])))
embeddings = embeddings[gt != 255]
gt = gt[gt != 255]
if "annotated" in data:
annotated_gt = data["annotated"]
results = {}
n_queries = {}
# sel_classes = [12, 22, 3, 34]
sel_classes = list(range(37))
for cl in sel_classes:
query_list = np.argwhere(gt == cl).flatten()
if "annotated" in data and query_list.size >= args["min_query_count"]:
query_list = np.array(
[q for q in query_list if annotated_gt[q] != 0])
n_queries[cl] = (len(query_list)
if len(query_list) >= args["min_query_count"] else 0)
if query_list.size >= args["min_query_count"]:
results[cl] = meanaverageprecision(
query_list,
gt,
embeddings,
distance_metric=args["distance_metric"],
gt_annotation=annotated_gt if "annotated" in data else None,
n_jobs=args["n_jobs"],
)
_log.info("{:>{width}s} ({:>4d}): {:>7.2%}".format(
trainid_to_name[cl],
len(query_list),
results[cl],
width=max([len(str(v)) for v in trainid_to_name.values()]),
))
_log.info("Average: {:.2%}".format(
sum(results.values()) / len(results.values())))
_log.info("Weighted Average: {:.2%}".format(
sum([v * n_queries[k]
for k, v in results.items()]) / sum(n_queries.values())))
if args["plot_dir"] is not None:
_log.info("Start plotting...")
fig = plt.figure(
"mAP values in % for retrieval in the embedding space")
ax = fig.add_subplot(111)
rects = ax.bar(
x=np.arange(len(results) + 2),
height=([v * 100 for k, v in results.items()] +
[sum(results.values()) / len(results.values()) * 100] + [
sum([v * n_queries[k] for k, v in results.items()]) /
sum(n_queries.values()) * 100
]),
)
ax.set_xticks(np.arange(len(results) + 2))
ax.set_xticklabels(labels=[trainid_to_name[k]
for k in results.keys()] + ["Average"] +
["Weighted Average"])
for rect in rects:
height = rect.get_height()
ax.annotate(
"{:.1f}".format(height),
xy=(rect.get_x() + rect.get_width() / 2, height),
xytext=(0, 3), # 3 points vertical offset
textcoords="offset points",
ha="center",
va="bottom",
)
# ax.title.set_text('Retrieval results in the embedding space')
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
ax.spines["bottom"].set_visible(False)
ax.yaxis.grid(True)
ax.xaxis.set_tick_params(rotation=50)
ax.set_ylabel("mAP in %")
ax.set_axisbelow(True)
plt.savefig(join(args["plot_dir"], "map_plot.eps"),
dpi=300,
bbox_inches="tight")
_log.info("Saved plot of mAP results to '{}'".format(
join(args["plot_dir"], "map_plot.eps")))
def main(args, tsne, _run, _log):
log_config(_run, _log)
with open(args['embeddings_file'], 'rb') as f:
data = pkl.load(f)
gt = np.array(data['gt']).squeeze()
_log.debug('Number of segments: {}'.format(gt.shape[0]))
gt = np.vectorize(id_to_trainid.get)(gt)
if (data['nn_embeddings'].shape[1] != args['embedding_size'] if 'nn_embeddings' in data.keys() else True)\
or args['overwrite_embeddings']:
embeddings = np.stack(data['embeddings'])
# _log.info('Standardizing embeddings...')
# embeddings = (embeddings - embeddings.mean()) / embeddings.std()
if args['embedding_size'] < embeddings.shape[1] if args[
'embedding_size'] is not None else False:
_log.info('Computing embeddings for nearest neighbor search...')
if args['method'] == 'TSNE':
_log.info(
'Using t-SNE with method \'{}\' and dimensionality {}'.
format(
'barnes_hut' if args['embedding_size'] < 4 else
'exact', args['embedding_size']))
embeddings = PCA(n_components=50 if args['embedding_size'] < 50
else 100).fit_transform(embeddings)
embeddings = TSNE(n_components=args['embedding_size'],
n_jobs=args['n_jobs'],
method='barnes_hut'
if args['embedding_size'] < 4 else 'exact',
**tsne).fit_transform(embeddings)
elif args['method'] == 'Isomap':
_log.info('Using Isomap method.')
embeddings = PCA(n_components=50 if args['embedding_size'] < 50
else 100).fit_transform(embeddings)
embeddings = Isomap(
n_components=args['embedding_size'],
n_jobs=args['n_jobs'],
).fit_transform(embeddings)
elif args['method'] == 'PCA':
_log.info('Using PCA method.')
embeddings = PCA(n_components=args['embedding_size']
).fit_transform(embeddings)
data['nn_embeddings'] = embeddings
_log.debug('Saving computed manifold to embeddings file.')
with open(args['embeddings_file'], 'wb') as f:
pkl.dump(data, f)
else:
_log.info('Leaving data as it is.')
else:
embeddings = data['nn_embeddings']
_log.info(
'Using precomputed embeddings ({} dimensions) for nearest neighbor search...'
.format(embeddings.shape[1]))
embeddings = embeddings[gt != 255]
gt = gt[gt != 255]
if 'annotated' in data:
annotated_gt = data['annotated']
results = {}
n_queries = {}
# sel_classes = [12, 22, 3, 34]
sel_classes = list(range(37))
for cl in sel_classes:
query_list = np.argwhere(gt == cl).flatten()
if 'annotated' in data and query_list.size >= args['min_query_count']:
query_list = np.array(
[q for q in query_list if annotated_gt[q] != 0])
n_queries[cl] = len(
query_list) if len(query_list) >= args['min_query_count'] else 0
if query_list.size >= args['min_query_count']:
results[cl] = meanaverageprecision(
query_list,
gt,
embeddings,
distance_metric=args['distance_metric'],
gt_annotation=annotated_gt if 'annotated' in data else None,
n_jobs=args['n_jobs'],
)
_log.info('{:>{width}s} ({:>4d}): {:>7.2%}'.format(
trainid_to_name[cl],
len(query_list),
results[cl],
width=max([len(str(v)) for v in trainid_to_name.values()])))
_log.info('Average: {:.2%}'.format(
sum(results.values()) / len(results.values())))
_log.info('Weighted Average: {:.2%}'.format(
sum([v * n_queries[k]
for k, v in results.items()]) / sum(n_queries.values())))
if args['plot_dir'] is not None:
_log.info('Start plotting...')
fig = plt.figure(
'mAP values in % for retrieval in the embedding space')
ax = fig.add_subplot(111)
rects = ax.bar(
x=np.arange(len(results) + 2),
height=([v * 100 for k, v in results.items()] +
[sum(results.values()) / len(results.values()) * 100] + [
sum([v * n_queries[k] for k, v in results.items()]) /
sum(n_queries.values()) * 100
]))
ax.set_xticks(np.arange(len(results) + 2))
ax.set_xticklabels(labels=[trainid_to_name[k]
for k in results.keys()] + ['Average'] +
['Weighted Average'])
for rect in rects:
height = rect.get_height()
ax.annotate(
'{:.1f}'.format(height),
xy=(rect.get_x() + rect.get_width() / 2, height),
xytext=(0, 3), # 3 points vertical offset
textcoords="offset points",
ha='center',
va='bottom')
# ax.title.set_text('Retrieval results in the embedding space')
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['bottom'].set_visible(False)
ax.yaxis.grid(True)
ax.xaxis.set_tick_params(rotation=50)
ax.set_ylabel('mAP in %')
ax.set_axisbelow(True)
plt.savefig(join(args['plot_dir'], 'map_plot.eps'),
dpi=300,
bbox_inches='tight')
_log.info('Saved plot of mAP results to \'{}\''.format(
join(args['plot_dir'], 'map_plot.eps')))
def main(args, _run, _log):
log_config(_run, _log)
if not args['only_plot']:
with open(args['embeddings_file'], 'rb') as f:
data = pkl.load(f)
image_indices = np.array(data['image_index'])
image_level_index = np.array(data['image_level_index'])
gt_segments = np.array(data['gt'])
boxes = np.array(data['box'])
inds = get_indices(
join(CONFIG.metaseg_io_path, 'input', 'deeplabv3plus', 'a2d2'))
if args['file_total_count'] is None:
total_num_instances = {cl: 0 for cl in id_to_trainid.keys()}
else:
with open(args['file_total_count'], 'rb') as f:
total_num_instances = pkl.load(f)
filtered_num_instances = {cl: 0 for cl in id_to_trainid.keys()}
for ind in tqdm.tqdm(inds):
pred, gt, img_path = probs_gt_load(ind,
join(CONFIG.metaseg_io_path,
'input', 'deeplabv3plus',
'a2d2'),
preds=True)
# count number of instances of each class of the minimum size in ground truth and prediction
for cl in np.unique(gt):
components_gt, counts_gt = label(gt == cl)
if args['file_total_count'] is None:
for c in range(1, counts_gt + 1):
segment_indices = np.argwhere(components_gt == c)
top, left = segment_indices.min(0)
bottom, right = segment_indices.max(0)
if (bottom - top) < args['min_height'] or (
right - left) < args['min_width']:
continue
else:
total_num_instances[cl] += 1
if ind in image_indices:
for b in boxes[(gt_segments == cl)
& (image_level_index == np.argwhere(
image_indices == ind).squeeze()), :]:
components_gt, instance_counts = return_and_update_instances(
components_gt, b)
filtered_num_instances[cl] += instance_counts
_log.info('Saving file with total counts...')
if args['file_total_count'] is None:
with open(args['save_file_total'], 'wb') as f:
pkl.dump(total_num_instances, f)
_log.info('Saving file with filtered counts...')
with open(args['save_file_filtered'], 'wb') as f:
pkl.dump(filtered_num_instances, f)
else:
with open(args['save_file_total'], 'rb') as f:
total_num_instances = pkl.load(f)
with open(args['save_file_filtered'], 'rb') as f:
filtered_num_instances = pkl.load(f)
_log.info('Start plotting')
# aggregate over training ids:
num_instances = {k: 0 for k in trainid_to_name.keys()}
f_num_instances = {k: 0 for k in trainid_to_name.keys()}
for k, v in total_num_instances.items():
num_instances[id_to_trainid[k]] += v
for k, v in filtered_num_instances.items():
f_num_instances[id_to_trainid[k]] += v
sel_classes = None
# sel_classes = [31, 22, 12, 34, 3, 35] # classes with many extracted instances
# sel_classes = [1, 4, 17, 24, 16, 18] # classes with few extracted instances
# start_angles = [45, 0, 10, 0, 0, 0]
start_angles = [0] * 6
fontsize = 8
fig = plt.figure('Class occurances filtered and not filtered',
figsize=(3.3, 2.5) if sel_classes is not None else
(10, 10),
dpi=args['dpi'])
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.serif'] = ['Times New Roman'
] + plt.rcParams['font.serif']
plt.rcParams['font.size'] = 6.0
def label_autopct(pct, allvals):
absolute = int(pct / 100.0 * np.sum(allvals))
return '{:.1f}%\n({:d})'.format(pct, absolute) if pct > 10 else ''
n = math.ceil(math.sqrt(len([1 for v in num_instances.values() if v > 0])))
cmap = plt.get_cmap('tab20c')
for i, k in enumerate([key for key, v in num_instances.items()
if v > 0] if sel_classes is None else sel_classes):
if num_instances[k] > 0:
ax = fig.add_subplot(n if sel_classes is None else 2,
n if sel_classes is None else 3, i + 1)
ax.text(
0.5,
1.0,
'{}'.format(trainid_to_name[k] if not trainid_to_name[k][-1].
isdigit() else trainid_to_name[k][:-2]),
horizontalalignment='center',
transform=ax.transAxes,
fontdict=dict(size=8),
)
ax.pie(
[num_instances[k] - f_num_instances[k], f_num_instances[k]],
radius=1.2,
colors=cmap(np.array([10, 5])),
startangle=start_angles[i] if sel_classes is not None else 0,
# autopct=lambda pct: '{:1.0f}%'.format(pct) if pct > 10 else '',
autopct=lambda pct: label_autopct(pct, [
num_instances[k] - f_num_instances[k], f_num_instances[k]
]),
pctdistance=0.65,
wedgeprops=dict(
width=1.0,
edgecolor='w',
linewidth=2,
),
textprops=dict(
# size=fontsize,
),
)
ax.set(aspect='equal')
fig.tight_layout(pad=0.0, h_pad=0.0, w_pad=0.6, rect=(0.0, 0.0, 1.0, 1.0))
plt.savefig(
join(
args['plot_dir'], 'instance_counts{}.{}'.format(
'' if sel_classes is None else '_selected',
args['plot_filetype'])),
dpi=args['dpi'],
)
_log.info('Saved instance counts plot to \'{}\''.format(
join(
args['plot_dir'], 'instance_counts{}.{}'.format(
'' if sel_classes is None else '_selected',
args['plot_filetype']))))
def main(args, _run, _log):
log_config(_run, _log)
# load a network architecture
_log.info('Loading {}...'.format(args['net']))
if args['net'] == 'vgg16':
net = feature_vgg16()
elif args['net'] == 'resnet18':
net = feature_resnet18()
elif args['net'] == 'resnet101':
net = feature_resnet101()
elif args['net'] == 'resnet152':
net = feature_resnet152()
elif args['net'] == 'wide_resnet101':
net = feature_wide_resnet101()
elif args['net'] == 'densenet201':
net = feature_densenet201()
else:
raise ValueError
net = net.cuda(args['gpu'])
net.eval()
# if no precomputed segments have been supplied, they have to be computed
if args['load_file'] is None:
_log.info('Loading Metrics...')
xa_all = []
start_others = []
pred_test = []
dataset_assignments = []
image_indices = []
# the first dataset of the 'datasets' configuration serves as source domain dataset. Metric statistics of this
# dataset are used to normalize the target domain metric statistics. This is why it has to get loaded too.
_log.info('{}...'.format(args['datasets'][0]))
xa, ya, x_names, class_names, xa_mean, xa_std, classes_mean, classes_std, *_, start, pred = load_data(
args['datasets'][0])
# Now load all other metric statistics and normalize them using the source domain mean and standard deviation
for i, d in enumerate(args['datasets'][1:], start=1):
_log.info('{} ...'.format(d))
num_imgs = get_indices(
join(CONFIG.metaseg_io_path, 'metrics', 'deeplabv3plus', d))
xa_tmp, *_, start_tmp, pred_tmp = load_data(
d,
num_imgs=num_imgs,
xa_mean=xa_mean,
xa_std=xa_std,
classes_mean=classes_mean,
classes_std=classes_std)
xa_all.append(xa_tmp)
pred_test.append(pred_tmp)
dataset_assignments += [i] * len(num_imgs)
image_indices += num_imgs
start_others.append(start_tmp)
# combine them into single arrays
xa_all = np.concatenate(xa_all).squeeze()
pred_test = np.concatenate(pred_test).squeeze()
dataset_assignments = np.array(dataset_assignments).squeeze()
image_indices = np.array(image_indices).squeeze()
for starts in start_others[1:]:
start_others[0] += [s + start_others[0][-1] for s in starts[1:]]
start_all = start_others[0]
del xa_tmp, start_tmp, pred_tmp, start_others
_log.debug('Shape of metrics array: {}'.format(xa_all.shape))
# Using the normalized metric statistics use a meta segmentation network pretrained on the source domain to
# predict IoU
_log.info('Predicting IoU...')
if args['meta_model'] == 'neural':
ya_pred_test = meta_nn_predict(
pretrained_model_path=args['meta_nn_path'],
x_test=xa_all,
gpu=args['gpu'])
elif args['meta_model'] == 'linear':
ya_pred_test, _ = regression_fit_and_predict(x_train=xa,
y_train=ya,
x_test=xa_all)
else:
raise ValueError('Meta model {} not supported.'.format(
args['meta_model']))
# This list will be used as an additional filter. Only segments with class predictions in this list will be
# picked for further processing.
pred_class_selection = [
# 0, # road
# 1, # sidewalk
# 2, # building
3, # wall
4, # fence
6, # traffic light
7, # traffic sign
# 8, # vegetation
# 9, # terrain
# 10, # sky
11, # person
12, # rider
13, # car
14, # truck
15, # bus
16, # train
17, # motorcycle
18, # bicycle
]
# Now the different filters are getting applied to the segments
_log.info('Filtering segments...')
inds = np.zeros(pred_test.shape[0]).astype(np.bool)
# Filter for the predicted IoU to be less than the supplied threshold
inds = np.logical_or(inds, (ya_pred_test < args['iou_threshold']))
# Filter for extracting segments with predefined class predictions
inds = np.logical_and(inds, np.isin(pred_test, pred_class_selection))
_log.info('Filtered components (not checked for minimum size):')
train_dat = getattr(
importlib.import_module(CONFIG.TRAIN_DATASET.module_name),
CONFIG.TRAIN_DATASET.class_name)(**CONFIG.TRAIN_DATASET.kwargs)
_log.info('\t{:^{width}s} | Filtered | Total'.format(
'Class name',
width=max([len(v[0]) for v in train_dat.pred_mapping.values()] +
[len('Class name')])))
for cl in np.unique(pred_test).flatten():
_log.info('\t{:^{width}s} | {:>8d} | {:<8d}'.format(
train_dat.pred_mapping[cl][0],
inds[pred_test == cl].sum(), (pred_test == cl).sum(),
width=max([len(v[0])
for v in train_dat.pred_mapping.values()] +
[len('Class name')])))
# Aggregating arguments for extraction of component information.
inds = np.argwhere(inds).flatten()
component_image_mapping = get_image_index_to_components(
inds, start_all)
p_args = [
(v, image_indices[k], ya_pred_test[start_all[k]:start_all[k + 1]],
args['datasets'][dataset_assignments[k]], args['min_height'],
args['min_width'], args['min_crop_height'],
args['min_crop_width'], 'deeplabv3plus')
for k, v in component_image_mapping.items()
]
# Extracting component information can be parallelized in a multiprocessing pool
_log.info('Extracting component information...')
with Pool(args['n_jobs']) as p:
r = list(
tqdm.tqdm(p.imap(wrapper_cutout_components, p_args),
total=len(p_args)))
r = [c for c in r if len(c['component_indices']) > 0]
_log.info('Computing embeddings...')
crops = {
'embeddings': [],
'image_path': [],
'image_index': [],
'component_index': [],
'box': [],
'gt': [],
'pred': [],
'dataset': [],
'model_name': [],
'image_level_index': [],
'iou_pred': []
}
# process all extracted crops and compute feature embeddings
for c in tqdm.tqdm(r):
# load image
preds, gt, image_path = probs_gt_load(
c['image_index'],
input_dir=join(CONFIG.metaseg_io_path, 'input',
c['model_name'], c['dataset']),
preds=True)
crops['image_path'].append(image_path)
crops['model_name'].append(c['model_name'])
crops['dataset'].append(c['dataset'])
crops['image_index'].append(c['image_index'])
crops['iou_pred'].append(c['iou_pred'])
image = Image.open(image_path).convert('RGB')
for i, b in enumerate(c['boxes']):
img = trans.ToTensor()(image.crop(b))
img = trans.Normalize(mean=imagenet_mean,
std=imagenet_std)(img)
crops['embeddings'].append(get_embedding(
img.unsqueeze(0), net))
crops['box'].append(b)
crops['component_index'].append(c['component_indices'][i])
crops['image_level_index'].append(len(crops['image_path']) - 1)
crops['gt'].append(
get_component_gt(gt, c['segment_indices'][i]))
crops['pred'].append(
get_component_pred(preds, c['segment_indices'][i]))
_log.info('Saving data...')
with open(args['save_file'], 'wb') as f:
pkl.dump(crops, f)
else:
with open(args['load_file'], 'rb') as f:
crops = pkl.load(f)
_log.info('Computing embeddings...')
boxes = np.array(crops['box']).squeeze()
image_level_index = np.array(crops['image_level_index']).squeeze()
crops['embeddings'] = []
for i, image_path in tqdm.tqdm(enumerate(crops['image_path']),
total=len(crops['image_path'])):
image = Image.open(image_path).convert('RGB')
for j in np.argwhere(image_level_index == i).flatten():
img = trans.ToTensor()(image.crop(boxes[j]))
img = trans.Normalize(mean=imagenet_mean,
std=imagenet_std)(img)
crops['embeddings'].append(get_embedding(
img.unsqueeze(0), net))
if 'plot_embeddings' in crops:
del crops['plot_embeddings']
if 'nn_embeddings' in crops:
del crops['nn_embeddings']
_log.info('Saving data...')
with open(args['save_file'], 'wb') as f:
pkl.dump(crops, f)
def train(args, _run, _log):
log_config(_run, _log)
os.makedirs(dirname(args["save_folder"]), exist_ok=True)
_log.info("Loading data...")
xa, ya, _, _, xa_mean, xa_std, classes_mean, classes_std, *_ = load_data(
args["dataset"])
xa_val, ya_val, *_ = load_data(
args["dataset_val"],
xa_mean=xa_mean,
xa_std=xa_std,
classes_mean=classes_mean,
classes_std=classes_std,
)
dat = MetricDataset([xa, ya])
dat_val = MetricDataset([xa_val, ya_val])
_log.info("Training dataset size: {}".format(len(dat)))
_log.info("Validation dataset size: {}".format(len(dat_val)))
datloader = DataLoader(dat,
args["batch_size"],
shuffle=True,
num_workers=args["n_jobs"])
valloader = DataLoader(dat_val,
args["batch_size"],
shuffle=True,
num_workers=args["n_jobs"])
_log.info("Initializing network...")
net = getattr(
importlib.import_module(
meta_models[args["meta_model_name"]].module_name),
meta_models[args["meta_model_name"]].class_name,
)(xa.shape[1],
**meta_models[args["meta_model_name"]].kwargs).cuda(args["gpu"])
optimizer = torch.optim.Adam(net.parameters(),
lr=args["learning_rate"],
weight_decay=args["weight_decay"])
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=[20, 40],
gamma=0.1)
crit = nn.BCEWithLogitsLoss().cuda(args["gpu"])
crit_val = nn.BCEWithLogitsLoss(reduction="none")
min_loss = float("inf")
for e in range(args["epochs"]):
_log.info("Epoch {}/{}".format(e + 1, args["epochs"]))
_log.info("Training phase...")
net.train()
avg_loss = []
for x, y in datloader:
optimizer.zero_grad()
x, y = x.cuda(args["gpu"]), y.cuda(args["gpu"])
out = net(x)
loss = crit(out, y)
loss.backward()
optimizer.step()
avg_loss.append(loss.item())
# _run.log_scalar('batch_loss', loss.item())
avg_loss = sum(avg_loss) / len(avg_loss)
_run.log_scalar("train_loss", avg_loss)
_log.info("Validation phase...")
net.eval()
avg_val_loss = []
with torch.no_grad():
for x, y in valloader:
x = x.cuda(args["gpu"])
out = net(x).data.cpu()
avg_val_loss.append(crit_val(out, y))
avg_val_loss = torch.cat(avg_val_loss).mean().item()
_run.log_scalar("val_loss", avg_val_loss)
if avg_val_loss < min_loss:
min_loss = avg_val_loss
_log.info("Average validation loss decreased, saved model.")
torch.save(
{
"state_dict": net.state_dict(),
"train_xa_mean": xa_mean,
"train_xa_std": xa_std,
"train_classes_mean": classes_mean,
"train_classes_std": classes_std,
},
join(args["save_folder"], args["net_name"]),
)
scheduler.step()
def main(args, _run, _log):
log_config(_run, _log)
# load a network architecture
_log.info("Loading {}...".format(args["net"]))
if args["net"] == "vgg16":
net = feature_vgg16()
elif args["net"] == "resnet18":
net = feature_resnet18()
elif args["net"] == "resnet101":
net = feature_resnet101()
elif args["net"] == "resnet152":
net = feature_resnet152()
elif args["net"] == "wide_resnet101":
net = feature_wide_resnet101()
elif args["net"] == "densenet201":
net = feature_densenet201()
else:
raise ValueError
net = net.cuda(args["gpu"])
net.eval()
# if no precomputed segments have been supplied, they have to be computed
if args["load_file"] is None:
_log.info("Loading Metrics...")
xa_all = []
start_others = []
pred_test = []
dataset_assignments = []
image_indices = []
# the first dataset of the 'datasets' configuration serves as source domain
# dataset. Metric statistics of this dataset are used to normalize the target
# domain metric statistics. This is why it has to get loaded too.
if args["meta_model"] == "neural" and all(
i in torch.load(args["meta_nn_path"]).keys()
for i in [
"train_xa_mean",
"train_xa_std",
"train_classes_mean",
"train_classes_std",
]
):
_log.info(
"Loading values for normalization from saved model file '{}'".format(
args["meta_nn_path"]
)
)
model_dict = torch.load(args["meta_nn_path"])
xa_mean = model_dict["train_xa_mean"]
xa_std = model_dict["train_xa_std"]
classes_mean = model_dict["train_classes_mean"]
classes_std = model_dict["train_classes_std"]
else:
_log.info("{}...".format(args["datasets"][0]))
(
xa,
ya,
x_names,
class_names,
xa_mean,
xa_std,
classes_mean,
classes_std,
*_,
start,
pred,
) = load_data(args["datasets"][0])
# Now load all other metric statistics and normalize them using the source
# domain mean and standard deviation
for i, d in enumerate(args["datasets"][1:], start=1):
_log.info("{} ...".format(d))
num_imgs = get_indices(
join(CONFIG.metaseg_io_path, "metrics", "deeplabv3plus", d)
)
xa_tmp, *_, start_tmp, pred_tmp = load_data(
d,
num_imgs=num_imgs,
xa_mean=xa_mean,
xa_std=xa_std,
classes_mean=classes_mean,
classes_std=classes_std,
)
xa_all.append(xa_tmp)
pred_test.append(pred_tmp)
dataset_assignments += [i] * len(num_imgs)
image_indices += num_imgs
start_others.append(start_tmp)
# combine them into single arrays
xa_all = np.concatenate(xa_all).squeeze()
pred_test = np.concatenate(pred_test).squeeze()
dataset_assignments = np.array(dataset_assignments).squeeze()
image_indices = np.array(image_indices).squeeze()
for starts in start_others[1:]:
start_others[0] += [s + start_others[0][-1] for s in starts[1:]]
start_all = start_others[0]
del xa_tmp, start_tmp, pred_tmp, start_others
_log.debug("Shape of metrics array: {}".format(xa_all.shape))
# Using the normalized metric statistics use a meta segmentation network
# pretrained on the source domain to predict IoU
_log.info("Predicting IoU...")
if args["meta_model"] == "neural":
ya_pred_test = meta_nn_predict(
pretrained_model_path=args["meta_nn_path"],
x_test=xa_all,
gpu=args["gpu"],
)
elif args["meta_model"] == "linear":
ya_pred_test, _ = regression_fit_and_predict(
x_train=xa, y_train=ya, x_test=xa_all
)
else:
raise ValueError("Meta model {} not supported.".format(args["meta_model"]))
# Now the different filters are getting applied to the segments
_log.info("Filtering segments...")
inds = np.zeros(pred_test.shape[0]).astype(np.bool)
# Filter for the predicted IoU to be less than the supplied threshold
inds = np.logical_or(inds, (ya_pred_test < args["iou_threshold"]))
# Filter for extracting segments with predefined class predictions
if hasattr(
importlib.import_module(CONFIG.TRAIN_DATASET.module_name),
"pred_class_selection",
):
pred_class_selection = getattr(
importlib.import_module(CONFIG.TRAIN_DATASET.module_name),
"pred_class_selection",
)
inds = np.logical_and(inds, np.isin(pred_test, pred_class_selection))
_log.info("Filtered components (not checked for minimum size):")
train_dat = getattr(
importlib.import_module(CONFIG.TRAIN_DATASET.module_name),
CONFIG.TRAIN_DATASET.class_name,
)(**CONFIG.TRAIN_DATASET.kwargs)
_log.info(
"\t{:^{width}s} | Filtered | Total".format(
"Class name",
width=max(
[len(v[0]) for v in train_dat.pred_mapping.values()]
+ [len("Class name")]
),
)
)
for cl in np.unique(pred_test).flatten():
_log.info(
"\t{:^{width}s} | {:>8d} | {:<8d}".format(
train_dat.pred_mapping[cl][0],
inds[pred_test == cl].sum(),
(pred_test == cl).sum(),
width=max(
[len(v[0]) for v in train_dat.pred_mapping.values()]
+ [len("Class name")]
),
)
)
# Aggregating arguments for extraction of component information.
inds = np.argwhere(inds).flatten()
component_image_mapping = get_image_index_to_components(inds, start_all)
p_args = [
(
v,
image_indices[k],
ya_pred_test[start_all[k] : start_all[k + 1]],
args["datasets"][dataset_assignments[k]],
args["min_height"],
args["min_width"],
args["min_crop_height"],
args["min_crop_width"],
"deeplabv3plus",
)
for k, v in component_image_mapping.items()
]
# Extracting component information can be parallelized in a multiprocessing pool
_log.info("Extracting component information...")
with Pool(args["n_jobs"]) as p:
r = list(
tqdm.tqdm(p.imap(wrapper_cutout_components, p_args), total=len(p_args))
)
r = [c for c in r if len(c["component_indices"]) > 0]
_log.info("Computing embeddings...")
crops = {
"embeddings": [],
"image_path": [],
"image_index": [],
"component_index": [],
"box": [],
"gt": [],
"pred": [],
"dataset": [],
"model_name": [],
"image_level_index": [],
"iou_pred": [],
}
# process all extracted crops and compute feature embeddings
for c in tqdm.tqdm(r):
# load image
preds, gt, image_path = probs_gt_load(
c["image_index"],
input_dir=join(
CONFIG.metaseg_io_path, "input", c["model_name"], c["dataset"]
),
preds=True,
)
crops["image_path"].append(image_path)
crops["model_name"].append(c["model_name"])
crops["dataset"].append(c["dataset"])
crops["image_index"].append(c["image_index"])
crops["iou_pred"].append(c["iou_pred"])
image = Image.open(image_path).convert("RGB")
for i, b in enumerate(c["boxes"]):
img = trans.ToTensor()(image.crop(b))
img = trans.Normalize(mean=imagenet_mean, std=imagenet_std)(img)
crops["embeddings"].append(get_embedding(img.unsqueeze(0), net))
crops["box"].append(b)
crops["component_index"].append(c["component_indices"][i])
crops["image_level_index"].append(len(crops["image_path"]) - 1)
crops["gt"].append(get_component_gt(gt, c["segment_indices"][i]))
crops["pred"].append(get_component_pred(preds, c["segment_indices"][i]))
_log.info("Saving data...")
with open(args["save_file"], "wb") as f:
pkl.dump(crops, f)
else:
with open(args["load_file"], "rb") as f:
crops = pkl.load(f)
_log.info("Computing embeddings...")
boxes = np.array(crops["box"]).squeeze()
image_level_index = np.array(crops["image_level_index"]).squeeze()
crops["embeddings"] = []
for i, image_path in tqdm.tqdm(
enumerate(crops["image_path"]), total=len(crops["image_path"])
):
image = Image.open(image_path).convert("RGB")
for j in np.argwhere(image_level_index == i).flatten():
img = trans.ToTensor()(image.crop(boxes[j]))
img = trans.Normalize(mean=imagenet_mean, std=imagenet_std)(img)
crops["embeddings"].append(get_embedding(img.unsqueeze(0), net))
if "plot_embeddings" in crops:
del crops["plot_embeddings"]
if "nn_embeddings" in crops:
del crops["nn_embeddings"]
_log.info("Saving data...")
with open(args["save_file"], "wb") as f:
pkl.dump(crops, f)