def test_cat_removal(get_data):
cas = COCO_Assistant(get_data[0], get_data[1])
test_ann = "tiny2.json"
test_rcats = sorted(['plane', 'ship', 'Large_Vehicle'])
cas.remove_cat(interactive=False, jc=test_ann, rcats=test_rcats)
orig = COCO(os.path.join(cas.ann_dir, cas.jc))
rmj = COCO(os.path.join(cas.resrm_dir, cas.jc))
orig_names = [
list(orig.cats.values())[i]['name'] for i in range(len(orig.cats))
]
rmj_names = [
list(rmj.cats.values())[i]['name'] for i in range(len(rmj.cats))
]
diff_names = sorted(list(set(orig_names) - set(rmj_names)))
# Clean up
shutil.rmtree(cas.resrm_dir)
if diff_names != test_rcats:
raise AssertionError(
"Failure in removing following categories: {}".format(test_rcats))
def test_merge(get_data):
cas = COCO_Assistant(get_data[0], get_data[1])
cas.merge()
comb = COCO(os.path.join(cas.resann_dir, 'merged.json'))
# Get combined annotation count
combann = len(comb.anns)
# Get individual annotation counts
ann_counts = [len(_cfile.anns) for _cfile in cas.annfiles]
print(combann)
print(sum(ann_counts))
# Clean up
shutil.rmtree(cas.res_dir)
if sum(ann_counts) != combann:
raise AssertionError("Failure in merging datasets")
def main(args):
cas = COCO_Assistant(args.img_dir, args.ann_dir)
if args.mode == "m+s" or args.mode == "m":
cas.merge()
if args.mode == "m+s":
cas.splitter(cas.dst_ann, args.split,
os.path.join(cas.res_dir, args.train),
os.path.join(cas.res_dir, args.test))
elif args.mode == "s":
cas.splitter(args.annotations, args.split, args.train, args.test)
else:
print('wrong mode')
# Rename images in annotation file
for i, js in enumerate(sorted(os.listdir(annotations))):
json_p = os.path.join(annotations, js)
with open(json_p) as j:
data = json.load(j)
# Access image dict
for im in data['images']:
# Change image name
im['file_name'] = im['file_name'][:-4] + "_d" + str(
i + 1) + im['file_name'][-4:]
with open(json_p[:-5] + "_modified" + ".json", 'w') as outfile:
json.dump(data, outfile)
# %%
# %%
from coco_assistant import COCO_Assistant
# Specify image and annotation directories
img_dir = os.path.join(os.getcwd(), 'dataset', 'images')
ann_dir = os.path.join(os.getcwd(), 'dataset', 'annotations')
print(img_dir)
# Create COCO_Assistant object
cas = COCO_Assistant(img_dir, ann_dir)
# %%
cas.merge(merge_images=True)
# %%
cas.visualise()
# %%