darwin_to_yolo

darwin_to_yolo.py

import os
import json
import argparse
import shutil
from PIL import Image
import cv2
import numpy as np

classes = {"rkn": 0, "fln": 1}

# Create an ArgumentParser
parser = argparse.ArgumentParser(description="Convert JSON files to TXT files")

# Add arguments for input and output directories
parser.add_argument("--input", required=True, help="Path to the directory containing JSON files")
parser.add_argument("--output", required=True, help="Path to the directory where TXT files will be saved")
parser.add_argument("--split", required=False, help="Split the dataset into train and val sets")
args = parser.parse_args()

# Use the provided input and output directory paths
source_folder = args.input
image_source = args.input + "/images"
annotation_darwin = source_folder + "/releases/latest/annotations"

target_folder = args.output
os.makedirs(target_folder, exist_ok=True)
labels_yolo = target_folder + "/labels"
os.makedirs(labels_yolo, exist_ok=True)
images_yolo = target_folder + "/images"
os.makedirs(images_yolo, exist_ok=True)

files = os.listdir(image_source)
for file in files:
    source_file = os.path.join(image_source, file)
    destination_file = os.path.join(images_yolo, file)
    shutil.copy2(source_file, destination_file) 

# Ensure output directory exists
for filename in os.listdir(annotation_darwin):  
    if filename.endswith('.json'):
        input_file_path = os.path.join(annotation_darwin, filename)
        with open(input_file_path, 'r') as json_file:
            try:
                data = json.load(json_file)
                no_ext = os.path.splitext(filename)[0]
                output_filename = no_ext + '.txt'
                output_file_path = os.path.join(labels_yolo, output_filename)
                with open(output_file_path, 'w') as txt_file: txt_file.write("")
                filename = no_ext + ".tif"
                if not os.path.exists(os.path.join(image_source, filename)):
                    filename = no_ext + ".png"
                with Image.open(os.path.join(image_source, filename)) as img:
                    width, height = img.size
                for ann in data["annotations"]:
                    cls = classes[ann["name"]]
                    segments = []
                    for e in  ann["polygon"]["paths"][0]:
                        #change x and y to be relative to the image size
                        x = max(0, min(e["x"] / width, 1))
                        y = max(0, min(e["y"] / height, 1))
                        segments.append(str(x) + " " + str(y))
                    stringtowrite = str(cls) + " " + " ".join(segments) + "\n"
                    print(stringtowrite)
                    with open(output_file_path, 'a') as txt_file:
                        txt_file.write(stringtowrite)
            except json.JSONDecodeError:
                print(f"Error decoding JSON in {filename}")

if args.split:
    if float(args.split) < 0 or float(args.split) > 1:
        print("Invalid split value. Split value must be between 0 and 1")
        exit()
    train_folder = target_folder + "/train"
    os.makedirs(train_folder, exist_ok=True)
    train_labels_yolo = train_folder + "/labels"
    os.makedirs(train_labels_yolo, exist_ok=True)
    train_images_yolo = train_folder + "/images"
    os.makedirs(train_images_yolo, exist_ok=True)

    val_folder = target_folder + "/val"
    os.makedirs(val_folder, exist_ok=True)
    val_labels_yolo = val_folder + "/labels"
    os.makedirs(val_labels_yolo, exist_ok=True)
    val_images_yolo = val_folder + "/images"
    os.makedirs(val_images_yolo, exist_ok=True)

    test_folder = target_folder + "/test"
    os.makedirs(test_folder, exist_ok=True)
    test_labels_yolo = test_folder + "/labels"
    os.makedirs(test_labels_yolo, exist_ok=True)
    test_images_yolo = test_folder + "/images"
    os.makedirs(test_images_yolo, exist_ok=True)

    files = os.listdir(labels_yolo)
    for file in files:
        if np.random.rand() < float(args.split):
            source_file = os.path.join(labels_yolo, file)
            destination_file = os.path.join(train_labels_yolo, file)
            shutil.copy2(source_file, destination_file) 
            source_file = os.path.join(images_yolo, file.replace(".txt", ".tif"))
            if not os.path.exists(source_file):
                source_file = os.path.join(images_yolo, file.replace(".txt", ".png"))
            destination_file = os.path.join(train_images_yolo, file.replace(".txt", ".tif"))
            if not os.path.exists(destination_file):
                destination_file = os.path.join(train_images_yolo, file.replace(".txt", ".png"))
            shutil.copy2(source_file, destination_file) 
        else:
            source_file = os.path.join(labels_yolo, file)
            destination_file = os.path.join(val_labels_yolo, file)
            shutil.copy2(source_file, destination_file) 
            source_file = os.path.join(images_yolo, file.replace(".txt", ".tif"))
            if not os.path.exists(source_file):
                source_file = os.path.join(images_yolo, file.replace(".txt", ".png"))
            destination_file = os.path.join(val_images_yolo, file.replace(".txt", ".tif"))
            if not os.path.exists(destination_file):
                destination_file = os.path.join(val_images_yolo, file.replace(".txt", ".png"))
            shutil.copy2(source_file, destination_file)

#read saved txt files and draw the bounding boxes
for filename in os.listdir(labels_yolo):
    if filename.endswith('.txt'):
        input_file_path = os.path.join(labels_yolo, filename)
        with open(input_file_path, 'r') as txt_file:
            try:
                no_ext = os.path.splitext(filename)[0]
                filename = no_ext + ".tif"
                if not os.path.exists(os.path.join(images_yolo, filename)):
                    filename = no_ext + ".png"
                img = cv2.imread(os.path.join(images_yolo, filename))
                for line in txt_file:
                    parts = line.split(" ")
                    cls = int(parts[0])
                    segments = []
                    for i in range(1, len(parts), 2):
                        x = int(float(parts[i]) * img.shape[1])
                        y = int(float(parts[i + 1]) * img.shape[0])
                        segments.append((x, y))
                    color = (0, 255, 0) if cls == 0 else (0, 0, 255)
                    cv2.polylines(img, [np.array(segments)], True, color, 2)
                cv2.imshow("image", img)
                if cv2.waitKey(0) == ord('q'):
                    break
            except json.JSONDecodeError:
                print(f"Error decoding JSON in {filename}")