Skip to content

cleanup code #29

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jun 8, 2023
Merged

cleanup code #29

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
829e2a5
cleanup code
GreatV Jun 8, 2023
4d0dbcc
fixed
GreatV Jun 8, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 46 additions & 41 deletions src/labelme2yolo/l2y.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,13 +106,13 @@ def extend_point_list(point_list, out_format="polygon"):
ymax = max(float(point) for point in point_list[1::2])

if out_format == "bbox":
x = xmin
y = ymin
w = xmax - xmin
h = ymax - ymin
x = x + w / 2
y = y + h / 2
return np.array([x, y, w, h])
x_i = xmin
y_i = ymin
w_i = xmax - xmin
h_i = ymax - ymin
x_i = x_i + w_i / 2
y_i = y_i + h_i / 2
return np.array([x_i, y_i, w_i, h_i])

return np.array([xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax])

Expand All @@ -123,12 +123,12 @@ def save_yolo_label(json_name, label_dir_path, target_dir, yolo_obj_list):
target_dir,
json_name.replace(".json", ".txt"))

with open(txt_path, "w+", encoding="utf-8") as f:
with open(txt_path, "w+", encoding="utf-8") as file:
for yolo_obj in yolo_obj_list:
label, points = yolo_obj
points = [str(item) for item in points]
yolo_obj_line = f"{label} {' '.join(points)}\n"
f.write(yolo_obj_line)
file.write(yolo_obj_line)


def save_yolo_image(json_data, json_path, image_dir_path, target_dir):
Expand Down Expand Up @@ -186,25 +186,26 @@ def _make_train_val_dir(self):

os.makedirs(yolo_path)

def _get_dataset_part_json_names(self, dataset_part: str):
'''Get json names in dataset_part folder'''
set_folder = os.path.join(self._json_dir, dataset_part)
json_names = []
for sample_name in os.listdir(set_folder):
set_dir = os.path.join(set_folder, sample_name)
if os.path.isdir(set_dir):
json_names.append(sample_name + '.json')
return json_names

def _train_test_split(self, folders, json_names, val_size, test_size):
'''Split json names to train, val, test'''
if (len(folders) > 0 and
'train' in folders and
'val' in folders and
'test' in folders):
train_folder = os.path.join(self._json_dir, 'train/')
train_json_names = [train_sample_name + '.json'
for train_sample_name in os.listdir(train_folder)
if os.path.isdir(os.path.join(train_folder, train_sample_name))] # noqa: E501

val_folder = os.path.join(self._json_dir, 'val/')
val_json_names = [val_sample_name + '.json'
for val_sample_name in os.listdir(val_folder)
if os.path.isdir(os.path.join(val_folder, val_sample_name))] # noqa: E501

test_folder = os.path.join(self._json_dir, 'test/')
test_json_names = [test_sample_name + '.json'
for test_sample_name in os.listdir(test_folder)
if os.path.isdir(os.path.join(test_folder, test_sample_name))] # noqa: E501
train_json_names = self._get_dataset_part_json_names('train')
val_json_names = self._get_dataset_part_json_names('val')
test_json_names = self._get_dataset_part_json_names('test')

return train_json_names, val_json_names, test_json_names

Expand Down Expand Up @@ -236,8 +237,9 @@ def convert(self, val_size, test_size):

# convert labelme object to yolo format object, and save them to files
# also get image from labelme json file and save them under images folder
for target_dir, json_names in zip(('train/', 'val/', 'test/'),
(train_json_names, val_json_names, test_json_names)): # noqa: E501
dirs = ('train/', 'val/', 'test/')
names = (train_json_names, val_json_names, test_json_names)
for target_dir, json_names in zip(dirs, names):

with Pool(NUM_THREADS) as pool:
for json_name in json_names:
Expand All @@ -252,8 +254,8 @@ def convert(self, val_size, test_size):
def covert_json_to_text(self, target_dir, json_name):
"""Convert json file to yolo format text file and save them to files"""
json_path = os.path.join(self._json_dir, json_name)
with open(json_path, encoding="utf-8") as f:
json_data = json.load(f)
with open(json_path, encoding="utf-8") as file:
json_data = json.load(file)

print(f"Converting {json_name} for {target_dir.replace('/', '')} ...")

Expand All @@ -268,10 +270,12 @@ def covert_json_to_text(self, target_dir, json_name):
yolo_obj_list)

def convert_one(self, json_name):
"""Convert one json file to yolo format text file and save them to files"""
json_path = os.path.join(self._json_dir, json_name)
json_data = json.load(open(json_path))
with open(json_path, encoding="utf-8") as file:
json_data = json.load(file)

print('Converting %s ...' % json_name)
print(f'Converting {json_name} ...')

img_path = save_yolo_image(json_data, json_name,
self._json_dir, '')
Expand Down Expand Up @@ -314,7 +318,7 @@ def _get_circle_shape_yolo_object(self, shape, img_h, img_w):
if shape['label'] in self._label_id_map:
label_id = self._label_id_map[shape['label']]
else:
print('label %s not in %s' % shape['label'], self._label_list)
print(f"label {shape['label']} not in {self._label_list}")

return label_id, yolo_center_x, yolo_center_y, yolo_w, yolo_h

Expand All @@ -333,26 +337,27 @@ def _get_other_shape_yolo_object(self, shape, img_h, img_w):
if shape['label'] in self._label_id_map:
label_id = self._label_id_map[shape['label']]
else:
print('label %s not in %s' % shape['label'], self._label_list)
print(f"label {shape['label']} not in {self._label_list}")

return label_id, points.tolist()

def _save_dataset_yaml(self):
yaml_path = os.path.join(
self._json_dir, 'YOLODataset/', 'dataset.yaml')

with open(yaml_path, 'w+') as yaml_file:
yaml_file.write('train: %s\n' %
os.path.join(self._image_dir_path, 'train/'))
yaml_file.write('val: %s\n' %
os.path.join(self._image_dir_path, 'val/'))
yaml_file.write('test: %s\n' %
os.path.join(self._image_dir_path, 'test/'))
yaml_file.write('nc: %i\n' % len(self._label_id_map))
with open(yaml_path, 'w+', encoding="utf-8") as yaml_file:

names_str = ''
train_dir = os.path.join(self._image_dir_path, 'train/')
val_dir = os.path.join(self._image_dir_path, 'val/')
test_dir = os.path.join(self._image_dir_path, 'test/')

names_str = ''
for label, _ in self._label_id_map.items():
names_str += "'%s', " % label
names_str += f"\"{label}\", "
names_str = names_str.rstrip(", ")
yaml_file.write("names: [%s]" % names_str)

content = (f"train: {train_dir}\nval: {val_dir}\ntest: {test_dir}\n"
f"nc: {len(self._label_id_map)}\n"
f"names: [{names_str}]")

yaml_file.write(content)