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() | |