SkiRacingImageProcessing/DataSetPrep.py at main · chpy04/SkiRacingImageProcessing · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import json
import os
import shutil
import random

# # Read the NDJSON file
# ndjson_file = 'exportedData.ndjson'
# json_file = 'data.json'

# # Collect JSON objects in a list
# data = []
# with open(ndjson_file, 'r') as file:
#     for line in file:
#         data.append(json.loads(line.strip()))

# # Write the list as a JSON array to a new file
# with open(json_file, 'w') as file:
#     json.dump(data, file, indent=4)

# print(f"Converted NDJSON to JSON and saved as {json_file}")

# Open and read the JSON file
with open('data.json', 'r') as file:
    data = json.load(file)

images_source_path = "jay1-others"
trainingPercentage = .9

output_path = "output"
train_output_path = os.path.join(output_path, "train")
images_train_path = os.path.join(train_output_path, "images")
labels_train_path = os.path.join(train_output_path, "labels")

val_output_path = os.path.join(output_path, "val")
images_val_path = os.path.join(val_output_path, "images")
labels_val_path = os.path.join(val_output_path, "labels")

os.makedirs(images_train_path, exist_ok=True)
os.makedirs(images_val_path, exist_ok=True)
os.makedirs(labels_train_path, exist_ok=True)
os.makedirs(labels_val_path, exist_ok=True)

yoloData = []

for image in data:
    name = image["data_row"]["external_id"]
    pixWidth = image["media_attributes"]["width"]
    pixHeight = image["media_attributes"]["height"]
    location = image["projects"]["cm5voz0cq02q707yicvga85fs"]["labels"][0]["annotations"]["objects"][0]["bounding_box"]
    pixX = location["left"] + (location["width"] / 2)
    pixY = location["top"] + (location["height"] / 2)

    x = pixX / pixWidth
    y = pixY / pixHeight
    w = location["width"] / pixWidth
    h = location["height"] / pixHeight

    boxStr = f"0 {x} {y} {w} {h}"

    yoloData.append({
        "name": name,
        "box": boxStr
    })

length = len(yoloData)

print(len(yoloData))

# Process each JSON object
for obj in yoloData:
    image_name = obj["name"]
    box_string = obj["box"]

    section = train_output_path if random.random() <= trainingPercentage else val_output_path

    # Copy the image to the Images folder
    source_image_path = os.path.join(images_source_path, image_name)
    dest_image_path = os.path.join(section, "images", image_name)

    if os.path.exists(source_image_path):
        shutil.copy(source_image_path, dest_image_path)
    else:
        print(f"Warning: {source_image_path} does not exist.")

    # Create a .txt file in the Labels folder with the box string
    label_file_name = os.path.splitext(image_name)[0] + ".txt"
    label_file_path = os.path.join(section, "labels", label_file_name)

    with open(label_file_path, "w") as label_file:
        label_file.write(box_string)

print("Processing complete. Files saved in the 'output' folder.")