File size: 3,216 Bytes
abd2a81 |
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 |
#!/usr/bin/env python3
###################################################################
# Use this script to convert .ply results from
# Li, M., Lafarge, F., Marlet, R.: Approximating shapes in images with low-complexitypolygons. In: CVPR (2020)
# to COCO .json format
###################################################################
import fnmatch
import os
import argparse
import skimage.io
import skimage.morphology
import numpy as np
from tqdm import tqdm
import json
from plyfile import PlyData, PlyElement
import shapely.geometry
import shapely.ops
from frame_field_learning import polygonize_utils, save_utils
def get_args():
argparser = argparse.ArgumentParser(description=__doc__)
argparser.add_argument(
'--ply_dirpath',
required=True,
type=str,
help='Path to the directory where the .ply are.')
argparser.add_argument(
'--mask_dirpath',
required=True,
type=str,
help='Path to the directory where the masks are (used to compute probability of each polygonal partition.')
argparser.add_argument(
'--output_filepath',
required=True,
type=str,
help='Filepath of the final .json.')
args = argparser.parse_args()
return args
def ply_to_json(ply_dirpath, mask_dirpath, output_filepath, mask_filename_format="{:012d}.png"):
filenames = fnmatch.filter(os.listdir(ply_dirpath), "*.ply")
all_annotations = []
for filename in tqdm(filenames, desc="Process ply files:"):
image_id = int(os.path.splitext(filename)[0])
# Load .ply
plydata = PlyData.read(os.path.join(ply_dirpath, filename))
x = plydata['vertex']['x']
y = 299 - plydata['vertex']['y']
pos = np.stack([x, y], axis=1)
vertex1 = plydata['edge'].data["vertex1"]
vertex2 = plydata['edge'].data["vertex2"]
edge_index = np.stack([vertex1, vertex2], axis=1)
edge = pos[edge_index]
linestrings = []
for e in edge:
linestrings.append(shapely.geometry.LineString(e))
# Load mask
mask_filename = mask_filename_format.format(image_id)
mask = 0 < skimage.io.imread(os.path.join(mask_dirpath, mask_filename))
# Convert to polygons
polygons = shapely.ops.polygonize(linestrings)
# Remove low prob polygons
filtered_polygons = []
filtered_polygon_probs = []
for polygon in polygons:
prob = polygonize_utils.compute_geom_prob(polygon, mask)
# print("simple:", np_indicator.min(), np_indicator.mean(), np_indicator.max(), prob)
if 0.5 < prob:
filtered_polygons.append(polygon)
filtered_polygon_probs.append(prob)
annotations = save_utils.poly_coco(filtered_polygons, filtered_polygon_probs, image_id)
all_annotations.extend(annotations)
with open(output_filepath, 'w') as outfile:
json.dump(all_annotations, outfile)
if __name__ == "__main__":
args = get_args()
ply_dirpath = args.ply_dirpath
mask_dirpath = args.mask_dirpath
output_filepath = args.output_filepath
ply_to_json(ply_dirpath, mask_dirpath, output_filepath)
|