86 lines
3.4 KiB
Python
86 lines
3.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
import json
|
|
import os
|
|
import argparse
|
|
from tqdm import tqdm
|
|
import chardet
|
|
|
|
# 預定義標籤類別與對應編號
|
|
PREDEFINED_CLASSES = {
|
|
"Edge bruises": 0
|
|
}
|
|
|
|
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, classes):
|
|
json_paths = [p for p in os.listdir(json_dir) if p.endswith('.json')]
|
|
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 in the same directory as JSON file
|
|
txt_path = os.path.join(json_dir, json_path.replace('.json', '.txt'))
|
|
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']
|
|
if shape_dict['shape_type'] == 'rectangle':
|
|
x1, y1 = points[0]
|
|
x2, y2 = points[1]
|
|
center_x = (x1 + x2) / 2 / w
|
|
center_y = (y1 + y2) / 2 / h
|
|
width = abs(x2 - x1) / w
|
|
height = abs(y2 - y1) / h
|
|
else:
|
|
print(f"Unsupported shape type: {shape_dict['shape_type']}")
|
|
continue
|
|
except (ValueError, KeyError) as e:
|
|
print(f"Error processing shape in file '{json_path}': {e}")
|
|
continue
|
|
|
|
label_str = f"{label_index} {center_x} {center_y} {width} {height}\n"
|
|
txt_file.writelines(label_str)
|
|
|
|
def extract_labels(json_dir):
|
|
labels = set()
|
|
json_paths = [p for p in os.listdir(json_dir) if p.endswith('.json')]
|
|
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='Convert JSON annotations to TXT format')
|
|
parser.add_argument('--json-dir', type=str, default='./img0204/P/', help='Directory containing JSON files')
|
|
args = parser.parse_args()
|
|
|
|
json_dir = args.json_dir
|
|
|
|
# Use predefined classes
|
|
predefined_classes = PREDEFINED_CLASSES
|
|
|
|
# Perform conversion
|
|
convert_label_json(json_dir, predefined_classes)
|