cvdatasets/scripts/display.py

#!/usr/bin/env python
if __name__ != '__main__': raise Exception("Do not import me!")
import sys
sys.path.insert(0, "..")

import logging
import numpy as np
import matplotlib.pyplot as plt

from argparse import ArgumentParser

from cvdatasets.annotation import AnnotationType
from utils import parser, plot_crops

def main(args):
	assert args.dataset in AnnotationType, \
		f"AnnotationType is not known: \"{args.dataset}\""

	annotation_cls = AnnotationType[args.dataset].value
	logging.info(f"Loading \"{args.dataset}\" annnotations from \"{args.data}\"")
	annot = annotation_cls(root_or_infofile=args.data, parts=args.parts, load_strict=False)

	kwargs = {}
	if annot.info is None:
		# features = args.features[0 if args.subset == "train" else 1]
		kwargs = dict(
			part_rescale_size=args.rescale_size,
			# features=features,
			uniform_parts=args.uniform_parts,
			ratio=args.ratio,
		)

	data = annot.new_dataset(
		args.subset,

		center_cropped=not args.no_center_crop,
		crop_to_bb=args.crop_to_bb,
		crop_uniform=args.crop_uniform,

		parts_in_bb=args.parts_in_bb,

		rnd_select=args.rnd,
		seed=args.seed,

		**kwargs
	)

	logging.info(f"Loaded {len(data)} {args.subset} images")

	if args.only_class >= 0:
		mask = data.labels == args.only_class
		logging.info(f"Showing only {mask.sum()} images from class {args.only_class}")
		idxs = np.where(mask)[0]
	else:
		start = max(args.start, 0)
		n_images = min(args.n_images, len(data) - start)
		end = max(start, start + n_images)
		logging.info(f"Showing only images {start} - {end}")
		idxs = range(start, end)

	for i in idxs:
		im, parts, label = data[i]

		fig1, axs = plt.subplots(1, 1, figsize=(16,9))
		axs = [axs]

		axs[0].axis("off")
		axs[0].set_title("Visible Parts")
		axs[0].imshow(im)

		if not args.crop_to_bb and not args.no_bboxes:
			data.plot_bounding_box(i, axs[0])

		# axs[1].axis("off")
		# axs[1].set_title("{}selected parts".format("randomly " if args.rnd else ""))
		# axs[1].imshow(parts.reveal(im, ratio=data.ratio))

		if not args.no_parts:
			parts.plot(im=im, ax=axs[0], ratio=data.ratio, linewidth=3)
			if data.uniform_parts:
				crop_names = None
			else:
				crop_names = list(data._annot.part_names.values())

			part_crops = parts.visible_crops(im, ratio=data.ratio)
			if args.rnd:
				parts.invert_selection()
				action_crops = parts.visible_crops(im, ratio=data.ratio)

			plot_crops(part_crops, f"{args.parts}: Selected parts", names=crop_names)

			if args.rnd:
				plot_crops(action_crops, f"{args.parts}: Actions", names=crop_names)

		plt.show()
		plt.close()

main(parser.parse_args())