Spaces:
Running
Running
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') |