-
Notifications
You must be signed in to change notification settings - Fork 0
/
dataset_transfer.py
47 lines (40 loc) · 2.05 KB
/
dataset_transfer.py
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
from os import makedirs, listdir
from os.path import exists
from shutil import copyfile
from cv2 import imwrite
from numpy import load, uint8
from rich.progress import Progress
def check_or_create(path: str) -> None:
if not exists(path):
makedirs(path)
def transfer(src: str, destination: str) -> None:
check_or_create(f"{destination}/nnUNet_preprocessed")
check_or_create(f"{destination}/nnUNet_weights")
destination = f"{destination}/nnUNet_raw/Dataset001_SegSTRONG_C"
check_or_create(destination)
with open(f"{destination}/dataset.json", "w") as f:
f.write("{\"channel_names\": {\"0\": \"red\", \"1\": \"green\", \"2\": \"blue\"}, "
"\"labels\": {\"background\": 0, \"tool\": 1}, \"numTraining\": 6600, \"file_ending\": \".png\"}")
check_or_create(f"{destination}/imagesTr")
check_or_create(f"{destination}/labelsTr")
with Progress() as progress:
task = progress.add_task("[white]Transferring...", total=3300)
serial = 0
src = f"{src}/train"
for i in listdir(src):
for j in listdir(f"{src}/{i}"):
for n in range(300):
copyfile(f"{src}/{i}/{j}/regular/left/{n}.png",
f"{destination}/imagesTr/case_{(s := str(serial).zfill(4))}_0000.png")
imwrite(f"{destination}/labelsTr/case_{s}.png",
load(f"{src}/{i}/{j}/ground_truth/left/{n}.npy").astype(uint8))
serial += 1
copyfile(f"{src}/{i}/{j}/regular/right/{n}.png",
f"{destination}/imagesTr/case_{(s := str(serial).zfill(4))}_0000.png")
imwrite(f"{destination}/labelsTr/case_{s}.png",
load(f"{src}/{i}/{j}/ground_truth/right/{n}.npy").astype(uint8))
serial += 1
progress.update(task, advance=1)
if __name__ == '__main__':
transfer("S:/SharedDatasets/SegSTRONGC_release/SegSTRONGC_release",
"S:/SharedDatasets/SegSTRONGC_release/SegSTRONGC_nnunet")