Spaces:
Runtime error
Runtime error
File size: 4,556 Bytes
6c06d1a |
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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 |
import argparse
import glob
import os
import os.path as osp
import cv2
import mmcv
import numpy as np
import pycocotools.mask as maskUtils
from mmengine.fileio import dump
from mmengine.utils import (Timer, mkdir_or_exist, track_parallel_progress,
track_progress)
def collect_files(img_dir, gt_dir):
files = []
img_files = glob.glob(osp.join(img_dir, 'image/*.tif'))
for img_file in img_files:
segm_file = gt_dir + '/label/' + os.path.basename(img_file)
files.append((img_file, segm_file))
assert len(files), f'No images found in {img_dir}'
print(f'Loaded {len(files)} images from {img_dir}')
return files
def collect_annotations(files, nproc=1):
print('Loading annotation images')
if nproc > 1:
images = track_parallel_progress(load_img_info, files, nproc=nproc)
else:
images = track_progress(load_img_info, files)
return images
def load_img_info(files):
img_file, segm_file = files
segm_img = mmcv.imread(segm_file, flag='unchanged', backend='cv2')
num_labels, instances, stats, centroids = cv2.connectedComponentsWithStats(segm_img, connectivity=4)
anno_info = []
for inst_id in range(1, num_labels):
category_id = 1
mask = np.asarray(instances == inst_id, dtype=np.uint8, order='F')
if mask.max() < 1:
print(f'Ignore empty instance: {inst_id} in {segm_file}')
continue
mask_rle = maskUtils.encode(mask[:, :, None])[0]
area = maskUtils.area(mask_rle)
# convert to COCO style XYWH format
bbox = maskUtils.toBbox(mask_rle)
# for json encoding
mask_rle['counts'] = mask_rle['counts'].decode()
anno = dict(
iscrowd=0,
category_id=category_id,
bbox=bbox.tolist(),
area=area.tolist(),
segmentation=mask_rle)
anno_info.append(anno)
video_name = osp.basename(osp.dirname(img_file))
img_info = dict(
# remove img_prefix for filename
file_name=osp.basename(img_file),
height=segm_img.shape[0],
width=segm_img.shape[1],
anno_info=anno_info,
segm_file=osp.basename(segm_file))
return img_info
def cvt_annotations(image_infos, out_json_name):
out_json = dict()
img_id = 0
ann_id = 0
out_json['images'] = []
out_json['categories'] = []
out_json['annotations'] = []
for image_info in image_infos:
image_info['id'] = img_id
anno_infos = image_info.pop('anno_info')
out_json['images'].append(image_info)
for anno_info in anno_infos:
anno_info['image_id'] = img_id
anno_info['id'] = ann_id
out_json['annotations'].append(anno_info)
ann_id += 1
img_id += 1
cat = dict(id=1, name='building')
out_json['categories'].append(cat)
if len(out_json['annotations']) == 0:
out_json.pop('annotations')
dump(out_json, out_json_name)
return out_json
def parse_args():
parser = argparse.ArgumentParser(
description='Convert WHU Building annotations to COCO format')
parser.add_argument('--cityscapes_path', default='/Users/kyanchen/datasets/Building/WHU', help='cityscapes data path')
parser.add_argument('--img-dir', default='', type=str)
parser.add_argument('--gt-dir', default='', type=str)
parser.add_argument('-o', '--out-dir', default='/Users/kyanchen/datasets/Building/WHU/annotations', help='output path')
parser.add_argument(
'--nproc', default=0, type=int, help='number of process')
args = parser.parse_args()
return args
def main():
args = parse_args()
cityscapes_path = args.cityscapes_path
out_dir = args.out_dir if args.out_dir else cityscapes_path
mkdir_or_exist(out_dir)
img_dir = osp.join(cityscapes_path, args.img_dir)
gt_dir = osp.join(cityscapes_path, args.gt_dir)
set_name = dict(
train='WHU_building_train.json',
val='WHU_building_val.json',
test='WHU_building_test.json'
)
for split, json_name in set_name.items():
print(f'Converting {split} into {json_name}')
with Timer(print_tmpl='It took {}s to convert Cityscapes annotation'):
files = collect_files(
osp.join(img_dir, split), osp.join(gt_dir, split))
image_infos = collect_annotations(files, nproc=args.nproc)
cvt_annotations(image_infos, osp.join(out_dir, json_name))
if __name__ == '__main__':
main()
|