Created
March 19, 2024 13:47
-
-
Save bresilla/dc4f8ec739e858bc79fbfa9073c0e07a to your computer and use it in GitHub Desktop.
darwin_to_yolo
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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}") |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment