File size: 1,736 Bytes
3f3265f | 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 | import os
import argparse
import numpy as np
import cv2
import json
from tqdm import tqdm
import pycocotools.mask as mask_util
parser = argparse.ArgumentParser()
parser.add_argument('--input', type=str, required=True)
parser.add_argument('--output', type=str, required=True)
args = parser.parse_args()
def rle_to_polygon(rle) -> (list, list, float):
mask = mask_util.decode(rle)
mask = np.where(mask > 0.5, 1, 0)
mask = np.ascontiguousarray(mask, dtype=np.uint8)
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)
polygon = []
segs = []
contours = [contour for contour in contours if contour.size >= 6]
for contour in contours:
_contour = contour.flatten().tolist()
_contour.append(_contour[0])
_contour.append(_contour[1])
polygon.extend(_contour)
segs.append(contour.flatten().tolist())
RLEs = mask_util.frPyObjects(segs, mask.shape[0], mask.shape[1])
RLE = mask_util.merge(RLEs)
area = float(mask_util.area(RLE))
return polygon, area
def main(args):
with open(args.input, 'r') as f:
ann_data = json.load(f)
for ann in tqdm(ann_data['annotations'], total=len(ann_data['annotations'])):
rle_data = ann['segmentation']
# encodedRLE = mask_util.frPyObjects(rle_data, rle_data['size'][0], rle_data['size'][1])
# polygon, area = rle_to_polygon(encodedRLE)
polygon, area = rle_to_polygon(rle_data)
ann['segmentation'] = [polygon]
ann['area'] = area
ann['iscrowd'] = 0
os.makedirs(os.path.dirname(args.output), exist_ok=True)
with open(args.output, 'w') as f:
json.dump(ann_data, f)
if __name__ == '__main__':
main(args)
|