-
Notifications
You must be signed in to change notification settings - Fork 11
/
Copy pathannotation.py
135 lines (124 loc) · 4.98 KB
/
annotation.py
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
import os
from PIL import Image
import numpy as np
from json import load, dump
from typing import List
class Object:
def __init__(self, category: str, group: int, segmentation, area, layer, bbox, iscrowd=0, note=''):
self.category = category
self.group = group
self.segmentation = segmentation
self.area = area
self.layer = layer
self.bbox = bbox
self.iscrowd = iscrowd
self.note = note
class Annotation:
def __init__(self, image_path, label_path, image=None, remote=False):
self.description = 'SAMLabeler'
self.remote = remote
if remote:
self.img_folder = ""
self.img_name = image_path
self.label_path = label_path
self.note = ''
if not isinstance(image, np.ndarray):
image = np.array(image)
else:
img_folder, img_name = os.path.split(image_path)
self.img_folder = img_folder
self.img_name = img_name
self.label_path = label_path
self.note = ''
if not isinstance(image, np.ndarray):
image = np.array(Image.open(image_path) if image is None else image)
if image.ndim == 3:
self.height, self.width, self.depth = image.shape
elif image.ndim == 2:
self.height, self.width = image.shape
self.depth = 0
else:
self.height, self.width, self.depth = image.shape[:, :3]
print('Warning: Except image has 2 or 3 ndim, but get {}.'.format(image.ndim))
del image
self.objects: List[Object, ...] = []
def load_from_dict(self, dataset):
info = dataset.get('info', {})
description = info.get('description', '')
if description in ['ISAT', 'SAMLabeler']:
# ISAT格式json
objects = dataset.get('objects', [])
self.img_name = info.get('name', '')
width = info.get('width', None)
if width is not None:
self.width = width
height = info.get('height', None)
if height is not None:
self.height = height
depth = info.get('depth', None)
if depth is not None:
self.depth = depth
self.note = info.get('note', '')
for obj in objects:
category = obj.get('category', 'unknow')
group = obj.get('group', 0)
if group is None: group = 0
segmentation = obj.get('segmentation', [])
iscrowd = obj.get('iscrowd', 0)
note = obj.get('note', '')
area = obj.get('area', 0)
layer = obj.get('layer', 2)
bbox = obj.get('bbox', [])
obj = Object(category, group, segmentation, area, layer, bbox, iscrowd, note)
self.objects.append(obj)
else:
# labelme格式json
shapes = dataset.get('shapes', {})
for shape in shapes:
# 只加载多边形
is_polygon = shape.get('shape_type', '') == 'polygon'
if not is_polygon:
continue
category = shape.get('label', 'unknow')
group = shape.get('group_id', 0)
if group is None: group = 0
segmentation = shape.get('points', [])
iscrowd = shape.get('iscrowd', 0)
note = shape.get('note', '')
area = shape.get('area', 0)
layer = shape.get('layer', 2)
bbox = shape.get('bbox', [])
obj = Object(category, group, segmentation, area, layer, bbox, iscrowd, note)
self.objects.append(obj)
def load_annotation(self):
if os.path.exists(self.label_path):
with open(self.label_path, 'r') as f:
dataset = load(f)
self.load_from_dict(dataset)
def to_dict(self):
dataset = {}
dataset['info'] = {}
dataset['info']['description'] = self.description
dataset['info']['folder'] = self.img_folder
dataset['info']['name'] = self.img_name
dataset['info']['width'] = self.width
dataset['info']['height'] = self.height
dataset['info']['depth'] = self.depth
dataset['info']['note'] = self.note
dataset['objects'] = []
for obj in self.objects:
object = {}
object['category'] = obj.category
object['group'] = obj.group
object['segmentation'] = obj.segmentation
object['area'] = obj.area
object['layer'] = obj.layer
object['bbox'] = obj.bbox
object['iscrowd'] = obj.iscrowd
object['note'] = obj.note
dataset['objects'].append(object)
return dataset
def save_annotation(self):
with open(self.label_path, 'w') as f:
dump(self.to_dict(), f) # , indent=4) # 节约空间
return True