File size: 3,885 Bytes
26364eb |
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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
import cv2
import json
import os
import numpy as np
from PIL import Image, ImageDraw
def convert_seg_coord_to_mask(img_size, coords):
"""Converts the segmentation coords found in COCO dataset to mask
Args:
img_size (_type_): _description_
coords (_type_): _description_
Returns:
_type_: _description_
"""
img = Image.new('L', img_size, 0)
for polygon in coords:
ImageDraw.Draw(img).polygon(polygon, outline=1, fill=1)
mask = np.array(img)
return mask
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']
self.task = cfg_obj['error_analysis']['task']
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])
img = Image.open(self.img_folder_path + img_fname)
for annot in self.annot_data['annotations']:
if img_id == annot['image_id']:
if annot['category_id'] in list(self.labels_dict.values()):
if self.task == "det":
annots.append([annot['category_id'],annot['bbox'][0],annot['bbox'][1],annot['bbox'][2],annot['bbox'][3]])
elif self.task == "seg":
# call convert_seg_coord_to_mask to convert segmentations [x1,y1,x2,y2,...,xn,yn] to binary mask
mask = convert_seg_coord_to_mask(img.size, annot['segmentation'])
annots.append([annot['category_id'], mask])
if self.task == "det":
return np.array(annots)
elif self.task == "seg":
return 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__':
import yaml
# get_annots()
cfg_file = open("cfg/cfg.yml")
cfg_obj = yaml.load(cfg_file, Loader=yaml.FullLoader)
annots_obj = AnnotsGTGetter(cfg_obj)
gt_dict = annots_obj.get_gt_annots()
# print (gt_dict)
# annots_obj.get_imgs()
# # to output mask
# img_annots = gt_dict['000000576031.jpg']
# import matplotlib.pyplot as plt
# plt.imshow(img_annots[0][1])
# plt.show()
|