def __call__(self, objects, image_fn):
basename = aitool.get_basename(image_fn)
root = ET.Element("annotations")
ET.SubElement(root, "folder").text = 'ROVOC'
ET.SubElement(root, "filename").text = image_fn
for data in objects:
obj = ET.SubElement(root, "object")
if self.classes is not None:
ET.SubElement(
obj, "name").text = self.classes[data['category_id'] - 1]
else:
ET.SubElement(obj, "name").text = str(data['category_id'])
ET.SubElement(obj, "type").text = "robndbox"
bndbox = ET.SubElement(obj, "robndbox")
thetaobb = data['thetaobb']
cx, cy, w, h, theta = thetaobb
ET.SubElement(bndbox, "cx").text = str(cx)
ET.SubElement(bndbox, "cy").text = str(cy)
ET.SubElement(bndbox, "w").text = str(w)
ET.SubElement(bndbox, "h").text = str(h)
ET.SubElement(bndbox, "angle").text = str(theta)
tree = ET.ElementTree(root)
tree.write(f"{self.output_dir}/{basename}.xml",
pretty_print=True,
xml_declaration=True,
encoding='utf-8')
def __call__(self, objects, image_fn):
basename = aitool.get_basename(image_fn)
root = ET.Element("annotations")
ET.SubElement(root, "folder").text = 'VOC'
ET.SubElement(root, "filename").text = image_fn
for data in objects:
obj = ET.SubElement(root, "object")
if self.classes is not None:
ET.SubElement(
obj, "name").text = self.classes[data['category_id'] - 1]
else:
ET.SubElement(obj, "name").text = str(data['category_id'])
ET.SubElement(obj, "type").text = "bndbox"
bndbox = ET.SubElement(obj, "bndbox")
bbox = data['bbox']
xmin, ymin, xmax, ymax = aitool.xywh2xyxy(bbox)
ET.SubElement(bndbox, "xmin").text = str(xmin)
ET.SubElement(bndbox, "ymin").text = str(ymin)
ET.SubElement(bndbox, "xmax").text = str(xmax)
ET.SubElement(bndbox, "ymax").text = str(ymax)
tree = ET.ElementTree(root)
tree.write(f"{self.output_dir}/{basename}.xml",
pretty_print=True,
xml_declaration=True,
encoding='utf-8')
def __init__(self,
ann_file,
classes=[''],
data_keys=['bbox', 'category_id', 'segmentation']):
"""parse coco annotation file
Args:
ann_file (str): coco annotation file
classes (list, optional): class ids. Defaults to [''].
data_keys (list, optional): parse which items. Defaults to ['bbox', 'category_id', 'segmentation'].
"""
self.data_keys = data_keys
self.coco = COCO(ann_file)
self.img_ids = self.coco.get_img_ids()
self.cat_ids = self.coco.get_cat_ids()
self.categories = self.coco.dataset['categories']
self.img_fns = []
print("begin to parse the coco annotation file")
self.objects, self.objects_with_id, self.img_name_with_id = dict(
), dict(), dict()
for img_id in tqdm.tqdm(self.img_ids):
img_info = self.coco.load_imgs([img_id])[0]
img_name = aitool.get_basename(img_info['file_name'])
self.img_fns.append(img_name)
ann_ids = self.coco.get_ann_ids(img_ids=[img_id])
ann_info = self.coco.load_anns(ann_ids)
self.objects[img_name] = self._convert_items(ann_info, img_info)
self.objects_with_id[img_id] = self.objects[img_name].copy()
self.img_name_with_id[img_name] = img_id
def _get_image_list(self, k_fold=0):
if self.imageset_file is not None:
print(
f"loading image list from imageset file: {self.imageset_file}")
raise NotImplementedError
else:
print(f"loading image list from image dir: {self.image_dir}")
random.seed(0)
image_file_list = aitool.get_file_list(self.image_dir,
self.image_format)
image_list = [
aitool.get_basename(image_file)
for image_file in image_file_list
]
image_list.sort()
random.shuffle(image_list)
if k_fold == 0:
if self.expand_image_list is None:
return image_list
else:
return image_list + self.expand_image_list
else:
splitted_image_list = self._split_image_list(image_list, k_fold)
return splitted_image_list
def __call__(self, objects, image_fn):
basename = aitool.get_basename(image_fn)
with open(f"{self.output_dir}/{basename}.txt", 'w') as f:
for data in objects:
if self.classes is not None:
object_info = [self.classes[data['category_id'] - 1]]
else:
object_info = [str(data['category_id'])]
pointobb = data['pointobb']
object_info.extend([str(_) for _ in pointobb] + ['\n'])
f.write(" ".join(object_info))
def __call__(self, objects, image_fn, team='CAPTAIN-VIPG-Plane'):
basename = aitool.get_basename(image_fn)
root = ET.Element("annotation")
source = ET.SubElement(root, "source")
ET.SubElement(source, "filename").text = basename + '.tif'
ET.SubElement(source, "origin").text = 'GF2/GF3'
research = ET.SubElement(root, "research")
ET.SubElement(research, "version").text = "4.0"
ET.SubElement(research, "provider").text = "Wuhan University"
ET.SubElement(research, "author").text = team
ET.SubElement(
research, "pluginname"
).text = "Airplane Detection and Recognition in Optical Images"
ET.SubElement(research, "pluginclass").text = "Detection"
ET.SubElement(research, "time").text = "2020-07-2020-11"
objects_handle = ET.SubElement(root, "objects")
for data in objects:
obj = ET.SubElement(objects_handle, "object")
ET.SubElement(obj, "coordinate").text = "pixel"
ET.SubElement(obj, "type").text = "rectangle"
ET.SubElement(obj, "description").text = "None"
possibleresult = ET.SubElement(obj, "possibleresult")
if self.classes is not None:
ET.SubElement(possibleresult,
"name").text = self.classes[data['category_id'] -
1]
else:
ET.SubElement(possibleresult,
"name").text = str(data['category_id'])
ET.SubElement(possibleresult,
"probability").text = str(data['score'])
points = ET.SubElement(obj, "points")
pointobb = data['pointobb']
for idx in range(5):
if idx == 4:
idx = 0
ET.SubElement(
points, "point"
).text = f"{pointobb[2 * idx]},{pointobb[2 * idx + 1]}"
tree = ET.ElementTree(root)
tree.write(f"{self.output_dir}/{basename}.xml",
pretty_print=True,
xml_declaration=True,
encoding='utf-8')
def __call__(self, objects, image_fn):
basename = aitool.get_basename(image_fn)
with open(f"{self.output_dir}/{basename}.txt", 'w') as f:
for data in objects:
if self.classes is not None:
object_info = [self.classes[data['category_id'] - 1]]
else:
object_info = [str(data['category_id'])]
bbox = data['bbox']
xmin, ymin, xmax, ymax = aitool.xywh2xyxy(bbox)
object_info.extend(
[str(xmin), str(ymin),
str(xmax), str(ymax)])
f.write(" ".join(object_info))
def _parse_expand_image_list(self, expand_image_list):
if isinstance(expand_image_list, list):
return expand_image_list
elif isinstance(expand_image_list, str):
return_list = []
with open(expand_image_list, 'r') as f:
lines = f.readlines()
for line in lines:
return_list.append(aitool.get_basename(line.strip()))
return return_list
elif expand_image_list is None:
return expand_image_list
else:
raise NotImplementedError(
f"don't support the type of expand_image_list: {type(expand_image_list)}"
)
def __call__(self, img, objects, ori_image_fn):
image_basename = aitool.get_basename(ori_image_fn)
for data in objects:
bbox = [int(_) for _ in data['bbox']]
bbox = self._fix_bbox_bound(img, bbox)
pointobb = data['pointobb']
x, y, w, h = bbox
if self._filter_invalid_bbox(bbox):
continue
img_save = img[y:y + h, x:x + w, :]
output_fn = [image_basename]
output_fn += [str(int(coord)) for coord in pointobb]
output_file = os.path.join(
self.output_dir, "_".join(output_fn) + self.output_format)
cv2.imwrite(output_file, img_save)
def __init__(self,
pkl_file,
ann_file,
keep_boundary=True,
score_threshold=0.05,
min_area=10):
self.keep_boundary = keep_boundary
self.score_threshold = score_threshold
self.min_area = min_area
if isinstance(pkl_file, str):
results = mmcv.load(pkl_file)
elif isinstance(pkl_file, (list, tuple)):
results = pkl_file
else:
raise TypeError(
f'do not support the pkl file type: {type(pkl_file)}')
coco = COCO(ann_file)
self.img_ids = coco.get_img_ids()
self.cat_ids = coco.get_cat_ids()
self.img_fns = []
self.img_sizes = []
print("begin to convert pkl file to specific format")
self.objects = dict()
self.objects_with_id = dict()
for idx, img_id in tqdm.tqdm(enumerate(self.img_ids)):
info = coco.load_imgs([img_id])[0]
img_name = aitool.get_basename(info['file_name'])
img_size = (info['height'], info['width'])
self.img_fns.append(img_name)
result = results[idx]
self.objects[img_name] = self._convert_items(result,
img_size=img_size)
self.objects_with_id[img_id] = self.objects[img_name].copy()
import aitool
if __name__ == "__main__":
pkl_file = '../mmdetv2-bc/results/plane/pl_v001_mask_rcnn_x101_32x4d_fpn_2x/pl_v001_mask_rcnn_x101_32x4d_fpn_2x_coco_results.pkl'
ann_file = './data/plane/v1/coco/annotations/plane_val.json'
img_dir = './data/plane/v1/val/images'
pkl_parser = aitool.PklParserMask(pkl_file, ann_file)
image_list = aitool.get_file_list(img_dir, '.tif')
for image_file in image_list:
basename = aitool.get_basename(image_file)
objects = pkl_parser(basename)
print("number of objects: ", len(objects))
