Spaces:
Sleeping
Sleeping
File size: 2,549 Bytes
b78b0dc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import cv2
import json
import os
import numpy as np
class AnnotsGTGetter:
def __init__(self, cfg_obj):
self.cfg_obj = cfg_obj
self.img_folder_path = cfg_obj['dataset']['img_folder_path']
self.json_folder_path = cfg_obj['dataset']['annotations_folder_path']
self.annot_json_fname = cfg_obj['dataset']['annotations_fname']
self.labels_dict = cfg_obj['error_analysis']['labels_dict']
json_file = open(self.json_folder_path + self.annot_json_fname)
self.annot_data = json.load(json_file)
self.img_ids_in_json = [annot['image_id'] for annot in self.annot_data['annotations']]
self.all_imgs = os.listdir(self.img_folder_path)
return
def get_imgs(self):
"""method to get the mutually -inclusive- images between the img_ids in json and those in the folder path
not needed because all images in folder were accounted for in the json...
"""
all_img_ids_in_folder = [int(i[:-4]) for i in self.all_imgs]
all_imgs_found = [i for i in all_img_ids_in_folder if i in self.img_ids_in_json]
print (len(all_imgs_found))
def get_annots(self, img_fname = '000000576052.jpg'):
"""retrieve annotation given a filename
Args:
img_fname (_type_): image file name
Returns:
np array: all annotations of an image
"""
# change img_fname for extraction purpose
# assumes jpg, png, but not jpeg...
# TODO - what if jpeg?
annots = []
img_id = int(img_fname[:-4])
for annot in self.annot_data['annotations']:
if img_id == annot['image_id']:
if annot['category_id'] in list(self.labels_dict.values()):
annots.append([annot['category_id'],annot['bbox'][0],annot['bbox'][1],annot['bbox'][2],annot['bbox'][3]])
return np.array(annots)
def get_gt_annots(self):
"""goes into the image folder, calls get_annots to extract image annotation
Returns:
dict: all annotations
"""
# create dictionary of gt annots
# for img in os.listdir(self.img_folder_path):
# self.get_annots(img)
all_gt_annots = {img: self.get_annots(img) for img in os.listdir(self.img_folder_path)}
return all_gt_annots
if __name__ == '__main__':
# get_annots()
annots_obj = AnnotsGTGetter()
gt_dict = annots_obj.get_gt_annots()
print (gt_dict)
# annots_obj.get_imgs()
|