File size: 1,178 Bytes
16a0f31 |
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 |
import os
import numpy as np
from sklearn.model_selection import train_test_split
import shutil
import argparse
from config import DATA_ROOT
dataset_root = os.path.join(DATA_ROOT, 'head_ct')
label_file = os.path.join(dataset_root, 'labels.csv')
data = np.loadtxt(label_file, dtype=int, delimiter=',', skiprows=1)
fnames = data[:, 0]
label = data[:, 1]
normal_fnames = fnames[label==0]
outlier_fnames = fnames[label==1]
target_root = '../datasets/HeadCT_anomaly_detection/headct'
train_root = os.path.join(target_root, 'train/good')
if not os.path.exists(train_root):
os.makedirs(train_root)
test_normal_root = os.path.join(target_root, 'test/good')
if not os.path.exists(test_normal_root):
os.makedirs(test_normal_root)
for f in normal_fnames:
source = os.path.join(dataset_root, 'head_ct/', '{:0>3d}.png'.format(f))
shutil.copy(source, test_normal_root)
test_outlier_root = os.path.join(target_root, 'test/defect')
if not os.path.exists(test_outlier_root):
os.makedirs(test_outlier_root)
for f in outlier_fnames:
source = os.path.join(dataset_root, 'head_ct/', '{:0>3d}.png'.format(f))
shutil.copy(source, test_outlier_root)
print('Done') |