Spaces:
Sleeping
Sleeping
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() | |