Created
March 19, 2024 13:47
-
-
Save bresilla/dc4f8ec739e858bc79fbfa9073c0e07a to your computer and use it in GitHub Desktop.
Revisions
-
bresilla created this gist
Mar 19, 2024 .There are no files selected for viewing
This file contains hidden or 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 charactersOriginal file line number Diff line number Diff line change @@ -0,0 +1,144 @@ 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}")