cvdatasets/tests/test_annotations.py

import unittest
import numpy as np
import os

from tests.configs import config
from os.path import *


from nabirds import BaseAnnotations, _MetaInfo
from nabirds.utils import read_info_file

class MockAnnotation(BaseAnnotations):
	name = "MOCK"

	@property
	def meta(self):
		info = _MetaInfo(
			images_folder="images",
			images_file="images.txt",
			labels_file="labels.txt",
			split_file="tr_ID.txt",
			bounding_boxes="bounding_boxes.txt",
			bounding_box_dtype=np.dtype([(v, np.int32) for v in "xywh"]),
			parts_file=join("parts", "part_locs.txt"),
			part_names_file=join("parts", "parts.txt"),
		)

		info.structure = [
			[info.images_file, "_images"],
			[info.labels_file, "labels"],
			[info.split_file, "_split"],
			[info.parts_file, "_part_locs"],
			[info.part_names_file, "_part_names"],
			[info.bounding_boxes, "_bounding_boxes"],
		]
		return info


class AnnotationTests(unittest.TestCase):

	def setUp(self):
		self.info = read_info_file(config.INFO_FILE)
		# self.annot = MockAnnotation()

	def tearDown(self):
		# clear mock data folder?
		pass

	def create_annotations(self, images, labels, split,
		bboxes=False,
		index_offset=0,
		n_parts=None):

		data_root = join(self.info.BASE_DIR, self.info.DATA_DIR)
		dataset_info = self.info.DATASETS.MOCK
		annot_dir = join(data_root, dataset_info.folder, dataset_info.annotations)

		if not isdir(annot_dir):
			os.makedirs(annot_dir)

		fname = lambda name: join(annot_dir, name)

		with open(fname("images.txt"), "w") as images_f,\
			open(fname("labels.txt"), "w") as labels_f,\
			open(fname("tr_ID.txt"), "w") as split_f:

			for i, (im, lab, sp) in enumerate(zip(images, labels, split), index_offset):
				print(i, im, file=images_f)
				print(lab, file=labels_f)
				print(sp, file=split_f)

		if bboxes:
			with open(fname("bounding_boxes.txt"), "w") as bbox_f:
				for i in range(index_offset, index_offset + len(images)):
					print(i, 0, 0, 100, 100, file=bbox_f)

		if n_parts is not None:
			parts_dir = join(annot_dir, "parts")
			if not isdir(parts_dir):
				os.makedirs(parts_dir)
			fname = lambda name: join(parts_dir, name)
			with open(fname("parts.txt"), "w") as part_names_f, \
				open(fname("part_locs.txt"), "w") as part_locs_f:

				for i in range(n_parts):
					print(i, "part_{}".format(i), file=part_names_f)

				for i in range(index_offset, len(images) + index_offset):
					for p in range(n_parts):
						print(i, 10*p, 10*p, 1, file=part_locs_f)


	def test_foo(self):
		self.create_annotations(
			["images{}.jpg".format(i) for i in range(10)],
			[i % 5 for i in range(10)],
			[int(i < 5) for i in range(10)],
			bboxes=False, n_parts=5
		)

		annot = MockAnnotation(config.INFO_FILE)

	def test_bar(self):
		self.assertFalse(0)