提交 612b40df 编写于 作者: K Kentaro Wada

Update labelme2voc.py and labelme2coco.py accordingly

上级 9a191df9
......@@ -8,6 +8,7 @@ import json
import os
import os.path as osp
import sys
import uuid
import numpy as np
import PIL.Image
......@@ -111,21 +112,28 @@ def main():
for shape in label_data['shapes']:
points = shape['points']
label = shape['label']
shape_type = shape.get('shape_type', None)
group_id = shape.get('group_id')
shape_type = shape.get('shape_type')
mask = labelme.utils.shape_to_mask(
img.shape[:2], points, shape_type
)
if label in masks:
masks[label] = masks[label] | mask
if group_id is None:
group_id = uuid.uuid1()
instance = (label, group_id)
if instance in masks:
masks[instance] = masks[instance] | mask
else:
masks[label] = mask
masks[instance] = mask
points = np.asarray(points).flatten().tolist()
segmentations[label].append(points)
segmentations[instance].append(points)
segmentations = dict(segmentations)
for label, mask in masks.items():
cls_name = label.split('-')[0]
for instance, mask in masks.items():
cls_name, group_id = instance
if cls_name not in class_name_to_id:
continue
cls_id = class_name_to_id[cls_name]
......@@ -139,7 +147,7 @@ def main():
id=len(data['annotations']),
image_id=image_id,
category_id=cls_id,
segmentation=segmentations[label],
segmentation=segmentations[instance],
area=area,
bbox=bbox,
iscrowd=0,
......
......@@ -103,7 +103,6 @@ def main():
img_shape=img.shape,
shapes=data['shapes'],
label_name_to_value=class_name_to_id,
type='instance',
)
ins[cls == -1] = 0 # ignore it.
......
......@@ -81,7 +81,7 @@ def main():
f.write(label_file.imageData)
img = labelme.utils.img_data_to_arr(label_file.imageData)
lbl = labelme.utils.shapes_to_label(
lbl, _ = labelme.utils.shapes_to_label(
img_shape=img.shape,
shapes=label_file.shapes,
label_name_to_value=class_name_to_id,
......
......@@ -41,7 +41,9 @@ def main():
else:
label_value = len(label_name_to_value)
label_name_to_value[label_name] = label_value
lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)
lbl, _ = utils.shapes_to_label(
img.shape, data['shapes'], label_name_to_value
)
label_names = [None] * (max(label_name_to_value.values()) + 1)
for name, value in label_name_to_value.items():
......
......@@ -50,7 +50,9 @@ def main():
else:
label_value = len(label_name_to_value)
label_name_to_value[label_name] = label_value
lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)
lbl, _ = utils.shapes_to_label(
img.shape, data['shapes'], label_name_to_value
)
label_names = [None] * (max(label_name_to_value.values()) + 1)
for name, value in label_name_to_value.items():
......
import math
import uuid
import numpy as np
import PIL.Image
......@@ -46,33 +47,31 @@ def shape_to_mask(img_shape, points, shape_type=None,
return mask
def shapes_to_label(img_shape, shapes, label_name_to_value, type='class'):
assert type in ['class', 'instance']
def shapes_to_label(img_shape, shapes, label_name_to_value):
cls = np.zeros(img_shape[:2], dtype=np.int32)
if type == 'instance':
ins = np.zeros(img_shape[:2], dtype=np.int32)
instance_names = ['_background_']
ins = np.zeros_like(cls)
instances = []
for shape in shapes:
points = shape['points']
label = shape['label']
group_id = shape.get('group_id')
if group_id is None:
group_id = uuid.uuid1()
shape_type = shape.get('shape_type', None)
if type == 'class':
cls_name = label
elif type == 'instance':
cls_name = label.split('-')[0]
if label not in instance_names:
instance_names.append(label)
ins_id = instance_names.index(label)
cls_name = label
instance = (cls_name, group_id)
if instance not in instances:
instances.append(instance)
ins_id = instances.index(instance) + 1
cls_id = label_name_to_value[cls_name]
mask = shape_to_mask(img_shape[:2], points, shape_type)
cls[mask] = cls_id
if type == 'instance':
ins[mask] = ins_id
ins[mask] = ins_id
if type == 'instance':
return cls, ins
return cls
return cls, ins
def labelme_shapes_to_label(img_shape, shapes):
......@@ -88,7 +87,7 @@ def labelme_shapes_to_label(img_shape, shapes):
label_value = len(label_name_to_value)
label_name_to_value[label_name] = label_value
lbl = shapes_to_label(img_shape, shapes, label_name_to_value)
lbl, _ = shapes_to_label(img_shape, shapes, label_name_to_value)
return lbl, label_name_to_value
......
......@@ -10,7 +10,7 @@ def test_shapes_to_label():
label_name = shape['label']
label_value = len(label_name_to_value)
label_name_to_value[label_name] = label_value
cls = shape_module.shapes_to_label(
cls, _ = shape_module.shapes_to_label(
img.shape, data['shapes'], label_name_to_value)
assert cls.shape == img.shape[:2]
......
......@@ -31,7 +31,7 @@ def get_img_and_lbl():
for label_name, label_value in label_name_to_value.items():
label_names[label_value] = label_name
lbl = shape_module.shapes_to_label(
lbl, _ = shape_module.shapes_to_label(
img.shape, data['shapes'], label_name_to_value
)
return img, lbl, label_names
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册