-
Notifications
You must be signed in to change notification settings - Fork 23
/
Copy pathpreprocessing.py
148 lines (131 loc) · 5.69 KB
/
preprocessing.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
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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
from __future__ import division
from configparser import ConfigParser
import argparse
from utils import extract_image_path, extract_n_preprocess_dicom, check_and_create_dir, extract_image, augment_image_pair
import imreg_dft as ird
import os
import cv2
from scipy.misc import imsave
from multiprocessing.pool import Pool
from itertools import product
import numpy as np
def similarity(params):
"""
Transform target image to make it similar to source image
"""
# The template
source_image, target_image, size, source_out_dir, target_out_dir, verbose = params
im0 = extract_n_preprocess_dicom(source_image, size)
# The image to be transformed
im1 = extract_n_preprocess_dicom(target_image, size)
filename = os.path.basename(os.path.normpath(source_image))
if verbose: print("Comparing %r .........." % filename)
# Transform im1 to im0
result = ird.similarity(im0, im1, numiter=3)
source_image_path = os.path.join(source_out_dir, filename)
target_image_path = os.path.join(target_out_dir, filename)
if verbose: print("Saving %r .........." % source_image_path)
imsave(source_image_path, im0)
imsave(target_image_path, result['timg'])
if verbose: print("Saved %r .........." % source_image_path)
def get_image_path_from(source_dir, target_dir):
"""
Get image paths from source and target directory
"""
source_images = extract_image_path([source_dir])
target_images = extract_image_path([target_dir])
assert len(source_images) == len(target_images), "Number of images in %r is not the same as %r" % (source_dir, target_dir)
return (source_images, target_images)
def check_n_create_output_dir(output_dir):
"""
Check and create output directory
"""
source_out_dir = os.path.join(output_dir, "source")
target_out_dir = os.path.join(output_dir, "target")
check_and_create_dir(source_out_dir)
check_and_create_dir(target_out_dir)
return (source_out_dir, target_out_dir)
def registration(verbose, num_threads, size, source_dir, target_dir, output_dir):
"""
Registrating images and save to output_dir
"""
if verbose: print("Get image paths ...")
# Get images paths
source_images, target_images = get_image_path_from(source_dir, target_dir)
# Check and create output directory
source_out_dir, target_out_dir = check_n_create_output_dir(output_dir)
pool = None
if num_threads >= 1:
pool = Pool(num_threads)
else:
pool = Pool()
job_args = [(source_images[i], target_images[i], size, source_out_dir, target_out_dir, verbose) for i in range(len(source_images))]
pool.map(similarity, job_args)
pool.close()
pool.join()
def augmentation_pair(params):
"""
Augmenting pair of images
"""
source_image_path, target_image_path, size, source_out_path, target_out_path, verbose = params
source_image = extract_image(source_image_path)
target_image = extract_image(target_image_path)
filename = os.path.basename(os.path.normpath(source_image_path))
if verbose:
print("Augmenting image %r ..." % filename)
augment_image_pair(source_image, target_image, size, source_out_path, target_out_path)
if verbose:
print("Augmented image %r ..." % filename)
def augmentation(verbose, num_threads, source_dir, target_dir, augmentation_seed, size, output_dir):
"""
Augment registered images and save to output_dir
"""
# Get images paths
if verbose: print("Get image paths ...")
source_images, target_images = get_image_path_from(source_dir, target_dir)
# Check and create output directory
source_out_dir, target_out_dir = check_n_create_output_dir(output_dir)
# Augmenting images
pool = None
if num_threads >= 1:
pool = Pool(num_threads)
else:
pool = Pool()
job_args = []
for seed in range(augmentation_seed):
for i in range(len(source_images)):
image_name = '%r_%r.png' % (seed, i)
source_out_path = os.path.join(source_out_dir, image_name)
target_out_path = os.path.join(target_out_dir, image_name)
job_args.append((source_images[i], target_images[i], size, source_out_path, target_out_path, verbose))
pool.map(augmentation_pair, job_args)
pool.close()
pool.join()
def main(args):
cp = ConfigParser()
cp.read(args.config)
verbose = cp["DATA"].getboolean("verbose")
num_threads = cp["DATA"].getint("num_threads")
image_size = cp["DATA"].getint("image_size")
# Data registration
source_dir = cp["REGISTRATION"].get("source_dir")
target_dir = cp["REGISTRATION"].get("target_dir")
output_dir = cp["REGISTRATION"].get("registered_output_dir")
is_registration = cp["REGISTRATION"].getboolean("data_registration")
if is_registration:
if verbose: print("Starting registration data ...")
registration(verbose, num_threads, image_size, source_dir, target_dir, output_dir)
# Data augmentation
source_dir = cp["AUGMENTATION"].get("source_dir")
target_dir = cp["AUGMENTATION"].get("target_dir")
output_dir = cp["AUGMENTATION"].get("augmented_output_dir")
augmentation_seed = cp["AUGMENTATION"].getint("augmentation_seed")
is_augmentation = cp["AUGMENTATION"].getboolean("data_augmentation")
if is_augmentation:
if verbose: print("Starting augmentation data ...")
augmentation(verbose, num_threads, source_dir, target_dir, augmentation_seed, image_size, output_dir)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='hmchuong - BoneSuppression v2 - Preprocessing data')
parser.add_argument('--config', default='config/data_preprocessing.cfg', type=str, help='config file')
args = parser.parse_args()
main(args)