96 lines
3.4 KiB
Python
96 lines
3.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
import json
|
|
import os
|
|
import argparse
|
|
from tqdm import tqdm
|
|
import chardet
|
|
|
|
# 預定義標籤類別與對應編號
|
|
PREDEFINED_CLASSES = {
|
|
"T7H": 0,
|
|
"T8H": 1,
|
|
"T10H": 2,
|
|
"T15H": 3,
|
|
"PH00": 4,
|
|
"PH0": 5,
|
|
"3/32": 6,
|
|
"1/8": 7
|
|
}
|
|
|
|
|
|
def read_json_file(file_path):
|
|
with open(file_path, "rb") as f:
|
|
raw_data = f.read()
|
|
detected_encoding = chardet.detect(raw_data)['encoding']
|
|
with open(file_path, "r", encoding=detected_encoding, errors="ignore") as f:
|
|
return json.load(f)
|
|
|
|
def convert_label_json(json_dir, save_dir, classes):
|
|
json_paths = os.listdir(json_dir)
|
|
os.makedirs(save_dir, exist_ok=True) # 確保目錄存在
|
|
for json_path in tqdm(json_paths, desc="Converting JSON files"):
|
|
json_path_full = os.path.join(json_dir, json_path)
|
|
try:
|
|
json_dict = read_json_file(json_path_full)
|
|
h, w = json_dict['imageHeight'], json_dict['imageWidth']
|
|
except (KeyError, json.JSONDecodeError) as e:
|
|
print(f"Error reading file '{json_path}': {e}")
|
|
continue
|
|
|
|
# Save txt path
|
|
txt_path = os.path.join(save_dir, json_path.replace('.json', '.txt'))
|
|
os.makedirs(os.path.dirname(txt_path), exist_ok=True)
|
|
with open(txt_path, 'w') as txt_file:
|
|
for shape_dict in json_dict['shapes']:
|
|
try:
|
|
label = shape_dict['label']
|
|
label_index = classes.get(label, None) # 查找標籤對應的編號
|
|
if label_index is None:
|
|
print(f"Label '{label}' not in predefined classes.")
|
|
continue
|
|
points = shape_dict['points']
|
|
except ValueError as e:
|
|
print(f"Error processing shape in file '{json_path}': {e}")
|
|
continue
|
|
|
|
points_nor_list = []
|
|
for point in points:
|
|
points_nor_list.append(point[0] / w)
|
|
points_nor_list.append(point[1] / h)
|
|
|
|
points_nor_list = list(map(lambda x: str(x), points_nor_list))
|
|
points_nor_str = ' '.join(points_nor_list)
|
|
|
|
label_str = str(label_index) + ' ' + points_nor_str + '\n'
|
|
txt_file.writelines(label_str)
|
|
|
|
def extract_labels(json_dir):
|
|
labels = set()
|
|
json_paths = os.listdir(json_dir)
|
|
print(f"Found {len(json_paths)} JSON files to process.")
|
|
for json_path in json_paths:
|
|
json_path_full = os.path.join(json_dir, json_path)
|
|
try:
|
|
json_dict = read_json_file(json_path_full)
|
|
for shape_dict in json_dict['shapes']:
|
|
labels.add(shape_dict['label'])
|
|
except (KeyError, json.JSONDecodeError) as e:
|
|
print(f"Error reading file '{json_path}': {e}")
|
|
print(f"Extracted labels: {labels}")
|
|
return list(labels)
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description='json convert to txt params')
|
|
parser.add_argument('--json-dir', type=str, default='E:/code/image_segmentation/labelme_json_dir/', help='json path dir')
|
|
parser.add_argument('--save-dir', type=str, default='E:/code/image_segmentation/labelme_txt_dir/', help='txt save dir')
|
|
args = parser.parse_args()
|
|
|
|
json_dir = args.json_dir
|
|
save_dir = args.save_dir
|
|
|
|
# 使用預定義類別
|
|
predefined_classes = PREDEFINED_CLASSES
|
|
|
|
# 進行轉換
|
|
convert_label_json(json_dir, save_dir, predefined_classes)
|