def setUp(self):
chainer.config.cv_read_image_backend = self.backend
self.file = tempfile.NamedTemporaryFile(suffix='.' + self.suffix,
delete=False)
self.path = self.file.name
if self.alpha is None:
if self.color:
self.img = np.random.randint(0,
255,
size=(3, ) + self.size,
dtype=np.uint8)
else:
self.img = np.random.randint(0,
255,
size=(1, ) + self.size,
dtype=np.uint8)
write_image(self.img, self.path)
else:
self.img = np.random.randint(0,
255,
size=(4, ) + self.size,
dtype=np.uint8)
_write_rgba_image(self.img, self.path)
def setUp(self):
if self.file_obj:
self.f = tempfile.TemporaryFile()
self.file = self.f
format = self.format
else:
if self.format == 'jpeg':
suffix = '.jpg'
else:
suffix = '.' + self.format
self.f = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
self.file = self.f.name
format = None
if self.alpha is None:
if self.color:
self.img = np.random.randint(0,
255,
size=(3, ) + self.size,
dtype=np.uint8)
else:
self.img = np.random.randint(0,
255,
size=(1, ) + self.size,
dtype=np.uint8)
write_image(self.img, self.file, format=format)
else:
self.img = np.random.randint(0,
255,
size=(4, ) + self.size,
dtype=np.uint8)
_write_rgba_image(self.img, self.file, format=format)
if self.file_obj:
self.file.seek(0)
def test_write_image(self):
if self.file_obj:
write_image(self.img, self.file, format=self.format)
self.file.seek(0)
else:
write_image(self.img, self.file)
img = Image.open(self.file)
W, H = img.size
self.assertEqual((H, W), self.size)
if self.color:
self.assertEqual(len(img.getbands()), 3)
else:
self.assertEqual(len(img.getbands()), 1)
if self.format in {'bmp', 'png'}:
img = np.asarray(img)
if img.ndim == 2:
# reshape (H, W) -> (1, H, W)
img = img[np.newaxis]
else:
# transpose (H, W, C) -> (C, H, W)
img = img.transpose((2, 0, 1))
np.testing.assert_equal(img, self.img)
def setUp(self):
self.temp_dir = tempfile.mkdtemp()
img_dir = os.path.join(self.temp_dir,
'leftImg8bit/{}/aachen'.format(self.split))
resol = 'gtFine' if self.label_mode == 'fine' else 'gtCoarse'
label_dir = os.path.join(self.temp_dir,
'{}/{}/aachen'.format(resol, self.split))
os.makedirs(img_dir)
os.makedirs(label_dir)
for i in range(10):
img = np.random.randint(0, 255,
size=(3, 128, 160)).astype(np.uint8)
write_image(
img,
os.path.join(
img_dir,
'aachen_000000_0000{:02d}_leftImg8bit.png'.format(i)))
label = np.random.randint(0, 34,
size=(1, 128, 160)).astype(np.int32)
write_image(
label,
os.path.join(
label_dir,
'aachen_000000_0000{:02d}_{}_labelIds.png'.format(
i, resol)))
self.dataset = CityscapesSemanticSegmentationDataset(
self.temp_dir, self.label_mode, self.split, self.ignore_labels)
def setUp(self):
self.file = tempfile.NamedTemporaryFile(
suffix='.' + self.suffix, delete=False)
self.path = self.file.name
self.img = np.random.randint(
0, 255, size=self.size, dtype=np.uint8)
write_image(self.img[np.newaxis], self.path)
def _save_img_file(path, size, color):
if color:
img = np.random.randint(
0, 255, size=(3,) + size, dtype=np.uint8)
else:
img = np.random.randint(
0, 255, size=(1,) + size, dtype=np.uint8)
write_image(img, path)
def _save_img_file(path, size, color):
if color:
img = np.random.randint(
0, 255, size=(3,) + size, dtype=np.uint8)
else:
img = np.random.randint(
0, 255, size=(1,) + size, dtype=np.uint8)
write_image(img, path)
def setUp(self):
self.file = tempfile.NamedTemporaryFile(
suffix='.' + self.suffix, delete=False)
self.path = self.file.name
if self.color:
self.img = np.random.randint(
0, 255, size=(3,) + self.size, dtype=np.uint8)
else:
self.img = np.random.randint(
0, 255, size=(1,) + self.size, dtype=np.uint8)
write_image(self.img, self.path)
def setUp(self):
self.file = tempfile.NamedTemporaryFile(
suffix='.' + self.suffix, delete=False)
self.path = self.file.name
if self.color:
self.img = np.random.randint(
0, 255, size=(3,) + self.size, dtype=np.uint8)
else:
self.img = np.random.randint(
0, 255, size=(1,) + self.size, dtype=np.uint8)
write_image(self.img, self.path)
def setUp(self):
self.temp_dir = tempfile.mkdtemp()
img_dir = os.path.join(self.temp_dir, 'leftImg8bit/test/berlin')
os.makedirs(img_dir)
for i in range(10):
img = np.random.randint(0, 255,
size=(3, 128, 160)).astype(np.uint8)
write_image(
img,
os.path.join(
img_dir,
'berlin_000000_0000{:02d}_leftImg8bit.png'.format(i)))
self.dataset = CityscapesTestImageDataset(self.temp_dir)
def setUp(self):
if self.file_obj:
self.f = tempfile.TemporaryFile()
self.file = self.f
format = self.format
else:
if self.format == 'jpeg':
suffix = '.jpg'
else:
suffix = '.' + self.format
self.f = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
self.file = self.f.name
format = None
self.img = np.random.randint(
0, 255, size=self.size, dtype=np.uint8)
write_image(self.img[None], self.file, format=format)
if self.file_obj:
self.file.seek(0)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', '-g', type=int, default=-1)
parser.add_argument('--pretrained-model')
parser.add_argument('--input-size', type=int, default=448)
args = parser.parse_args()
label_names = voc_semantic_segmentation_label_names
colors = voc_semantic_segmentation_label_colors
n_class = len(label_names)
input_size = (args.input_size, args.input_size)
model = get_pspnet_resnet50(n_class)
chainer.serializers.load_npz(args.pretrained_model, model)
if args.gpu >= 0:
chainer.cuda.get_device_from_id(args.gpu).use()
model.to_gpu(args.gpu)
dataset = get_sbd_augmented_voc()
for i in range(1, 100):
img = dataset[i][0]
# img = read_image(args.image)
labels = model.predict([img])
label = labels[0]
from chainercv.utils import write_image
write_image(label[None], '{}.png'.format(i))
fig = plt.figure()
ax1 = fig.add_subplot(1, 2, 1)
vis_image(img, ax=ax1)
ax2 = fig.add_subplot(1, 2, 2)
ax2, legend_handles = vis_semantic_segmentation(img,
label,
label_names,
colors,
ax=ax2)
ax2.legend(handles=legend_handles, bbox_to_anchor=(1, 1), loc=2)
plt.show()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', '-g', type=int, default=-1)
parser.add_argument('--pretrained-model')
parser.add_argument('--input-size', type=int, default=448)
args = parser.parse_args()
label_names = voc_semantic_segmentation_label_names
colors = voc_semantic_segmentation_label_colors
n_class = len(label_names)
input_size = (args.input_size, args.input_size)
model = get_pspnet_resnet50(n_class)
chainer.serializers.load_npz(args.pretrained_model, model)
if args.gpu >= 0:
chainer.cuda.get_device_from_id(args.gpu).use()
model.to_gpu(args.gpu)
dataset = get_sbd_augmented_voc()
for i in range(1, 100):
img = dataset[i][0]
# img = read_image(args.image)
labels = model.predict([img])
label = labels[0]
from chainercv.utils import write_image
write_image(
label[None], '{}.png'.format(i))
fig = plt.figure()
ax1 = fig.add_subplot(1, 2, 1)
vis_image(img, ax=ax1)
ax2 = fig.add_subplot(1, 2, 2)
ax2, legend_handles = vis_semantic_segmentation(
img, label, label_names, colors, ax=ax2)
ax2.legend(handles=legend_handles, bbox_to_anchor=(1, 1), loc=2)
plt.show()
def main():
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--model_path',
default='./progressive_growing_of_gans/Gs_chainer.npz')
args = parser.parse_args()
chainer.config.train = False
latent = np.random.randn(4, 512).astype(np.float32)
generator = Generator()
chainer.serializers.load_npz(args.model_path, generator)
with chainer.no_backprop_mode():
img = generator(latent)
print(img.shape)
# [-1, 1] -> [0, 255]
image = cuda.to_cpu(img.array) * 127.5 + 127.5
image = image.clip(0.0, 255.0).astype(np.float32)
utils.write_image(utils.tile_images(image, 2), 'out.png')
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', type=int, default=-1)
parser.add_argument('--batchsize', type=int, default=6)
parser.add_argument('--pretrained-model', type=str, default=None)
parser.add_argument('--out', type=str, default='result')
args = parser.parse_args()
# Model
model = get_pspnet_resnet50(len(voc_semantic_segmentation_label_names))
if args.gpu >= 0:
# Make a specified GPU current
chainer.cuda.get_device_from_id(args.gpu).use()
model.to_gpu() # Copy the model to the GPU
if args.pretrained_model:
chainer.serializers.load_npz(args.pretrained_model, model)
raw_train_data = get_sbd_augmented_voc()
train_data = raw_train_data
debug = True
for i in range(10):
if i < 6:
lr = 0.001
else:
lr = 0.0001
out = os.path.join(args.out, 'epoch_{0:02d}'.format(i))
train_one_epoch(model, train_data, lr, args.gpu, args.batchsize, out)
print('finished training a epoch')
labels = predict_all(model, raw_train_data)
if debug:
for i in range(len(labels)):
label = labels[i]
write_image(label[None],
os.path.join(out, 'image_{0:05d}.png'.format(i)))
train_data = TupleDataset(raw_train_data, labels)
def test_write_image(self):
write_image(self.img, self.path)
img = Image.open(self.path)
W, H = img.size
self.assertEqual((H, W), self.size)
if self.color:
self.assertEqual(len(img.getbands()), 3)
else:
self.assertEqual(len(img.getbands()), 1)
if self.suffix in {'bmp', 'png'}:
img = np.asarray(img)
if img.ndim == 2:
# reshape (H, W) -> (1, H, W)
img = img[np.newaxis]
else:
# transpose (H, W, C) -> (C, H, W)
img = img.transpose((2, 0, 1))
np.testing.assert_equal(img, self.img)
def test_write_image(self):
write_image(self.img, self.path)
img = Image.open(self.path)
W, H = img.size
self.assertEqual((H, W), self.size)
if self.color:
self.assertEqual(len(img.getbands()), 3)
else:
self.assertEqual(len(img.getbands()), 1)
if self.suffix in {'bmp', 'png'}:
img = np.asarray(img)
if img.ndim == 2:
# reshape (H, W) -> (1, H, W)
img = img[np.newaxis]
else:
# transpose (H, W, C) -> (C, H, W)
img = img.transpose(2, 0, 1)
np.testing.assert_equal(img, self.img)
h, w = dataset.crop
# converted image
if args.imgtype == "dcm":
if dname != prevdir:
salt = str(random.randint(1000, 999999))
prevdir = dname
for j in range(args.num_slices):
ref_dicom = dataset.overwrite(new[j], cnt, salt)
path = os.path.join(
outdir, '{:s}_{}_{}.dcm'.format(fn, args.suffix, j))
ref_dicom.save_as(path)
else:
path = os.path.join(outdir,
'{:s}_{}.jpg'.format(fn, args.suffix))
write_image(new, path)
## images for analysis
if args.output_analysis:
# original
path = os.path.join(outdir, '{:s}_0orig.png'.format(fn))
write_image((imgs[i] * 127.5 + 127.5).astype(np.uint8), path)
# cycle
path = os.path.join(outdir, '{:s}_1cycle.png'.format(fn))
write_image((cycle[i] * 127.5 + 127.5).astype(np.uint8), path)
# cycle difference
path = os.path.join(outdir, '{:s}_2cycle_diff.png'.format(fn))
# cycle_diff[i] = (cycle_diff[i]+1)/(imgs[i]+2) # [0,2]/[1,3] = (0.0,1.5)
cycle_diff[i] = np.abs(0.5 * cycle_diff[i])
print("cycle diff: {} {} {}".format(np.min(cycle_diff[i]),
np.mean(cycle_diff[i]),
out_v = gen(imgs)
if args.gpu >= 0:
imgs = xp.asnumpy(imgs.data)
out = xp.asnumpy(out_v.data)
else:
imgs = imgs.data
out = out_v.data
## output images
for i in range(len(out)):
fn = dataset.get_img_path(cnt)
print("\nProcessing {}".format(fn))
new = dataset.var2img(out[i])
print("raw value: {} {}".format(np.min(out[i]),np.max(out[i])))
path = os.path.join(outdir,os.path.basename(fn))
# converted image
if args.imgtype=="dcm":
ref_dicom = dataset.overwrite(new[0],fn,salt)
ref_dicom.save_as(path)
else:
write_image(new, path)
cnt += 1
####
elapsed_time = time.time() - start
print ("{} images in {} sec".format(cnt,elapsed_time))
# converted image
if args.imgtype == "dcm":
if os.path.dirname(fn) != prevdir:
salt = str(random.randint(1000, 999999))
prevdir = os.path.dirname(fn)
ref_dicom = dataset.overwrite(new[0], fn, salt)
path = os.path.join(
outdir, '{:s}_{}.dcm'.format(os.path.basename(fn),
args.suffix))
ref_dicom.save_as(path)
ch, cw = ref_dicom.pixel_array.shape
else:
path = os.path.join(
outdir, '{:s}_{}.jpg'.format(os.path.basename(fn),
args.suffix))
write_image(new, path)
## images for analysis
if args.output_analysis:
# original
path = os.path.join(
outdir, '{:s}_0org.jpg'.format(os.path.basename(fn)))
write_image((imgs[i] * 127.5 + 127.5).astype(np.uint8), path)
# converted
path = os.path.join(
outdir, '{:s}_3out.jpg'.format(os.path.basename(fn)))
write_image((out[i] * 127.5 + 127.5).astype(np.uint8), path)
# cycle
path = os.path.join(
outdir, '{:s}_1cycle.jpg'.format(os.path.basename(fn)))
write_image((cycle[i] * 127.5 + 127.5).astype(np.uint8), path)
os.makedirs(os.path.join(args.out,"trainA"), exist_ok=True)
os.makedirs(os.path.join(args.out,"trainB"), exist_ok=True)
for fullname in sorted(glob.glob(os.path.join(args.root,"**/*.{}".format(args.imgtype)), recursive=True)):
fn = os.path.basename(fullname)
fn,ext = os.path.splitext(fn)
if args.imgtype == 'dcm':
subdirname = os.path.basename(os.path.dirname(fullname))
ref_dicom_in = dicom.read_file(fullname, force=True)
ref_dicom_in.file_meta.TransferSyntaxUID = dicom.uid.ImplicitVRLittleEndian
dt=ref_dicom_in.pixel_array.dtype
fileB = "trainB/{}_{}_clean.dcm".format(subdirname,fn)
ref_dicom_in.save_as(os.path.join(args.out,fileB))
dat = ref_dicom_in.pixel_array
print(np.min(dat),np.max(dat))
# noise
dat = (ref_dicom_in.pixel_array + np.random.poisson(args.noise,ref_dicom_in.pixel_array.shape)).astype(dt)
print(np.min(dat),np.max(dat))
ref_dicom_in.PixelData = dat.tostring()
fileA = "trainA/{}_{}_noise.dcm".format(subdirname,fn)
ref_dicom_in.save_as(os.path.join(args.out,fileA))
else:
dat = read_image(fullname)
c,h,w = dat.shape
fileB = "trainB/{}_clean.jpg".format(fn)
write_image(dat,os.path.join(args.out,fileB))
dat += np.random.poisson(args.noise,dat.shape)
fileA = "trainA/{}_noise.jpg".format(fn)
write_image(dat,os.path.join(args.out,fileA))
print("{}\t{}".format(fileA,fileB))
if args.load_dataset is None:
data_dir = root
else:
data_dir = os.path.join(root, args.load_dataset)
data_dir = os.path.join(data_dir, "train{}".format(args.category.upper()))
dataset = Dataset(path=data_dir,
resize_to=args.resize_to,
crop_to=args.crop_to,
flip=False)
iterator = chainer.iterators.SerialIterator(dataset,
args.batch_size,
repeat=False,
shuffle=False)
xp = gen.xp
cnt = 0
for batch in iterator:
imgs = chainer.dataset.concat_examples(batch, device=args.gpu)
with chainer.using_config('train', False):
out = xp.asnumpy(gen(imgs).data)
for i in range(len(out)):
path = '{:s}/{:s}.jpg'.format(args.out, dataset.ids[cnt])
arr = (out[i] + 1.0) / 2.0 * 255.0
org = read_image(dataset.get_img_path(cnt))
_, h, w = org.shape
arr = resize(arr, (h, w))
write_image(arr, path)
cnt += 1
print(cnt)
def update_core(self):
optimizer_sd = self.get_optimizer('main')
optimizer_enc = self.get_optimizer('enc')
optimizer_dec = self.get_optimizer('dec')
optimizer_dis = self.get_optimizer('dis')
xp = self.seed.xp
step = self.iteration % self.args.iter
osem_step = step % self.args.osem
if step == 0:
batch = self.get_iterator('main').next()
self.prImg, self.rev, self.patient_id, self.slice = self.converter(batch, self.device)
print(self.prImg.shape)
self.n_reconst += 1
self.recon_freq = 1
if ".npy" in self.args.model_image:
self.seed.W.array = xp.reshape(xp.load(self.args.model_image),(1,1,self.args.crop_height,self.args.crop_width))
elif ".dcm" in self.args.model_image:
ref_dicom = dicom.read_file(self.args.model_image, force=True)
img = xp.array(ref_dicom.pixel_array+ref_dicom.RescaleIntercept)
img = (2*(xp.clip(img,self.args.HU_base,self.args.HU_base+self.args.HU_range)-self.args.HU_base)/self.args.HU_range-1.0).astype(np.float32)
self.seed.W.array = xp.reshape(img,(1,1,self.args.crop_height,self.args.crop_width))
else:
# initializers.Uniform(scale=0.5)(self.seed.W.array)
initializers.HeNormal()(self.seed.W.array)
self.initial_seed = self.seed.W.array.copy()
# print(xp.min(self.initial_seed),xp.max(self.initial_seed),xp.mean(self.initial_seed))
## for seed array
arr = self.seed()
HU = self.var2HU(arr)
raw = self.HU2raw(HU)
self.seed.cleargrads()
loss_seed = Variable(xp.array([0.0],dtype=np.float32))
# conjugate correction using system matrix
if self.args.lambda_sd > 0:
self.seed.W.grad = xp.zeros_like(self.seed.W.array)
loss_sd = 0
for i in range(len(self.prImg)):
if self.rev[i]:
rec_sd = F.exp(-F.sparse_matmul(self.prMats[osem_step],F.reshape(raw[i,:,::-1,::-1],(-1,1)))) ##
else:
rec_sd = F.exp(-F.sparse_matmul(self.prMats[osem_step],F.reshape(raw[i],(-1,1)))) ##
if self.args.log:
loss_sd += F.mean_squared_error(F.log(rec_sd),F.log(self.prImg[i][osem_step]))
else:
loss_sd += F.mean_squared_error(rec_sd,self.prImg[i][osem_step])
if self.args.system_matrix:
gd = F.sparse_matmul( self.conjMats[osem_step], rec_sd-self.prImg[i][osem_step], transa=True)
if self.rev[i]:
self.seed.W.grad[i] -= self.args.lambda_sd * F.reshape(gd, (1,self.args.crop_height,self.args.crop_width)).array[:,::-1,::-1] # / logrep.shape[0] ?
else:
self.seed.W.grad[i] -= self.args.lambda_sd * F.reshape(gd, (1,self.args.crop_height,self.args.crop_width)).array # / logrep.shape[0] ?
if not self.args.system_matrix:
(self.args.lambda_sd *loss_sd).backward()
chainer.report({'loss_sd': loss_sd/len(self.prImg)}, self.seed)
if self.args.lambda_tvs > 0:
loss_tvs = losses.total_variation(arr, tau=self.args.tv_tau, method=self.args.tv_method)
loss_seed += self.args.lambda_tvs * loss_tvs
chainer.report({'loss_tvs': loss_tvs}, self.seed)
if self.args.lambda_advs>0:
L_advs = F.average( (self.dis(arr)-1.0)**2 )
loss_seed += self.args.lambda_advs * L_advs
chainer.report({'loss_advs': L_advs}, self.seed)
## generator output
arr_n = losses.add_noise(arr,self.args.noise_gen)
if self.args.no_train_seed:
arr_n.unchain()
if not self.args.decoder_only:
arr_n = self.encoder(arr_n)
gen = self.decoder(arr_n) # range = [-1,1]
## generator loss
loss_gen = Variable(xp.array([0.0],dtype=np.float32))
plan, plan_ae = None, None
if self.args.lambda_ae1>0 or self.args.lambda_ae2>0:
plan = losses.add_noise(Variable(self.converter(self.get_iterator('planct').next(), self.device)), self.args.noise_dis)
plan_enc = self.encoder(plan)
plan_ae = self.decoder(plan_enc)
loss_ae1 = F.mean_absolute_error(plan,plan_ae)
loss_ae2 = F.mean_squared_error(plan,plan_ae)
if self.args.lambda_reg>0:
loss_reg_ae = losses.loss_func_reg(plan_enc[-1],'l2')
chainer.report({'loss_reg_ae': loss_reg_ae}, self.seed)
loss_gen += self.args.lambda_reg * loss_reg_ae
loss_gen += self.args.lambda_ae1 * loss_ae1 + self.args.lambda_ae2 * loss_ae2
chainer.report({'loss_ae1': loss_ae1}, self.seed)
chainer.report({'loss_ae2': loss_ae2}, self.seed)
if self.args.lambda_tv > 0:
L_tv = losses.total_variation(gen, tau=self.args.tv_tau, method=self.args.tv_method)
loss_gen += self.args.lambda_tv * L_tv
chainer.report({'loss_tv': L_tv}, self.seed)
if self.args.lambda_adv>0:
L_adv = F.average( (self.dis(gen)-1.0)**2 )
loss_gen += self.args.lambda_adv * L_adv
chainer.report({'loss_adv': L_adv}, self.seed)
## regularisation on the latent space
if self.args.lambda_reg>0:
loss_reg = losses.loss_func_reg(arr_n[-1],'l2')
chainer.report({'loss_reg': loss_reg}, self.seed)
loss_gen += self.args.lambda_reg * loss_reg
self.encoder.cleargrads()
self.decoder.cleargrads()
loss_gen.backward()
loss_seed.backward()
chainer.report({'loss_gen': loss_gen}, self.seed)
optimizer_enc.update()
optimizer_dec.update()
optimizer_sd.update()
chainer.report({'grad_sd': F.average(F.absolute(self.seed.W.grad))}, self.seed)
if hasattr(self.decoder, 'latent_fc'):
chainer.report({'grad_gen': F.average(F.absolute(self.decoder.latent_fc.W.grad))}, self.seed)
# reconstruction consistency for NN
if (step % self.recon_freq == 0) and self.args.lambda_nn>0:
self.encoder.cleargrads()
self.decoder.cleargrads()
self.seed.cleargrads()
gen.grad = xp.zeros_like(gen.array)
HU_nn = self.var2HU(gen)
raw_nn = self.HU2raw(HU_nn)
loss_nn = 0
for i in range(len(self.prImg)):
if self.rev[i]:
rec_nn = F.exp(-F.sparse_matmul(self.prMats[osem_step],F.reshape(raw_nn[i,:,::-1,::-1],(-1,1))))
else:
rec_nn = F.exp(-F.sparse_matmul(self.prMats[osem_step],F.reshape(raw_nn[i],(-1,1))))
loss_nn += F.mean_squared_error(rec_nn,self.prImg[i][osem_step])
if self.args.system_matrix:
gd_nn = F.sparse_matmul( rec_nn-self.prImg[i][osem_step], self.conjMats[osem_step], transa=True )
if self.rev[i]:
gen.grad[i] -= self.args.lambda_nn * F.reshape(gd_nn, (1,self.args.crop_height,self.args.crop_width)).array[:,::-1,::-1]
else:
gen.grad[i] -= self.args.lambda_nn * F.reshape(gd_nn, (1,self.args.crop_height,self.args.crop_width)).array
chainer.report({'loss_nn': loss_nn/len(self.prImg)}, self.seed)
if self.args.system_matrix:
gen.backward()
else:
(self.args.lambda_nn * loss_nn).backward()
if not self.args.no_train_seed:
optimizer_sd.update()
if not self.args.no_train_enc:
optimizer_enc.update()
if not self.args.no_train_dec:
optimizer_dec.update()
if self.seed.W.grad is not None:
chainer.report({'grad_sd_consistency': F.average(F.absolute(self.seed.W.grad))}, self.seed)
if hasattr(self.decoder, 'latent_fc'):
chainer.report({'grad_gen_consistency': F.average(F.absolute(self.decoder.latent_fc.W.grad))}, self.seed)
elif hasattr(self.decoder, 'ul'):
chainer.report({'grad_gen_consistency': F.average(F.absolute(self.decoder.ul.c1.c.W.grad))}, self.seed)
chainer.report({'seed_diff': F.mean_absolute_error(self.initial_seed,self.seed.W)/F.mean_absolute_error(self.initial_seed,xp.zeros_like(self.initial_seed))}, self.seed)
# clip seed to [-1,1]
if self.args.clip:
self.seed.W.array = xp.clip(self.seed.W.array,a_min=-1.0, a_max=1.0)
# adjust consistency loss update frequency
self.recon_freq = max(1,int(round(self.args.max_reconst_freq * (step-self.args.reconst_freq_decay_start) / (self.args.iter+1-self.args.reconst_freq_decay_start))))
## for discriminator
fake = None
if self.args.dis_freq > 0 and ( (step+1) % self.args.dis_freq == 0) and (self.args.lambda_gan+self.args.lambda_adv+self.args.lambda_advs>0):
# get mini-batch
if plan is None:
plan = self.converter(self.get_iterator('planct').next(), self.device)
plan = losses.add_noise(Variable(plan),self.args.noise_dis)
# create fake
if self.args.lambda_gan>0:
if self.args.decoder_only:
fake_seed = xp.random.uniform(-1,1,(1,self.args.latent_dim)).astype(np.float32)
else:
fake_seed = self.encoder(xp.random.uniform(-1,1,(1,1,self.args.crop_height,self.args.crop_width)).astype(np.float32))
fake = self.decoder(fake_seed)
# decoder
self.decoder.cleargrads()
loss_gan = F.average( (self.dis(fake)-1.0)**2 )
chainer.report({'loss_gan': loss_gan}, self.seed)
loss_gan *= self.args.lambda_gan
loss_gan.backward()
optimizer_dec.update(loss=loss_gan)
fake_copy = self._buffer.query(fake.array)
if self.args.lambda_nn>0:
fake_copy = self._buffer.query(self.converter(self.get_iterator('mvct').next(), self.device))
if (step+1) % (self.args.iter // 30):
fake_copy = Variable(self._buffer.query(gen.array))
# discriminator
L_real = F.average( (self.dis(plan)-1.0)**2 )
L_fake = F.average( self.dis(fake_copy)**2 )
loss_dis = 0.5*(L_real+L_fake)
self.dis.cleargrads()
loss_dis.backward()
optimizer_dis.update()
chainer.report({'loss_dis': (L_real+L_fake)/2}, self.seed)
if ((self.iteration+1) % self.args.vis_freq == 0) or ((step+1)==self.args.iter):
for i in range(self.args.batchsize):
outlist=[]
if not self.args.no_train_seed and not self.args.decoder_only:
outlist.append((self.seed()[i],"0sd"))
if plan_ae is not None:
outlist.append((plan[i],'2pl'))
outlist.append((plan_ae[i],'3ae'))
if self.args.lambda_nn>0 or self.args.lambda_adv>0:
if self.args.decoder_only:
gen_img = self.decoder([self.seed()])
else:
gen_img = self.decoder(self.encoder(self.seed()))
outlist.append((gen_img[i],'1gn'))
if fake is not None:
outlist.append((fake[i],'4fa'))
for out,typ in outlist:
out.to_cpu()
HU = (((out+1)/2 * self.args.HU_range)+self.args.HU_base).array # [-1000=air,0=water,>1000=bone]
print("type: ",typ,"HU:",np.min(HU),np.mean(HU),np.max(HU))
#visimg = np.clip((out.array+1)/2,0,1) * 255.0
b,r = -self.args.HU_range_vis//2,self.args.HU_range_vis
visimg = (np.clip(HU,b,b+r)-b)/r * 255.0
fn = 'n{:0>5}_iter{:0>6}_p{}_z{}_{}'.format(self.n_reconst,step+1,self.patient_id[i],self.slice[i],typ)
write_image(np.uint8(visimg),os.path.join(self.args.out,fn+'.jpg'))
if (step+1)==self.args.iter or (not self.args.no_save_dcm):
#np.save(os.path.join(self.args.out,fn+'.npy'),HU[0])
write_dicom(os.path.join(self.args.out,fn+'.dcm'),HU[0])
args = parser.parse_args()
print(args)
if args.gpu >= 0:
chainer.cuda.get_device_from_id(args.gpu).use()
gen = getattr(net, args.gen_class)()
if args.load_gen_model != '':
serializers.load_npz(args.load_gen_model, gen)
print('Generator model loaded')
if args.gpu >= 0:
gen.to_gpu()
print('use gpu {}'.format(args.gpu))
xp = gen.xp
img = read_image(args.input)
img = img.astype('f')
img = img * 2 / 255.0 - 1.0 # [-1, 1)
height, width = img.shape[1:]
img = np.expand_dims(img, axis=0)
img = xp.asarray(img)
with chainer.using_config('train', False):
out = gen(img)
out = resize(xp.asnumpy(out.data[0]), (height, width))
out = (out + 1.0) / 2.0 * 255.0
write_image(out, args.output)
def main():
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', type=int, default=-1)
parser.add_argument(
'--model',
default='',
help='if not specified, you download and use a pre-trained model.')
parser.add_argument('--snapshot', default='')
parser.add_argument('--image', type=str)
args = parser.parse_args()
if not args.image:
ValueError('args.image should be specified.')
else:
args.image = os.path.expanduser(args.image)
detector = SSD512(pretrained_model='voc0712')
model = StackedHG(16)
if args.model:
chainer.serializers.load_npz(args.model, model)
elif args.snapshot:
chainer.serializers.load_npz(snap2model_trainer(args.snapshot), model)
else:
# pre-trained model
model_path = './models/model_2018_05_22.npz'
if not os.path.exists(model_path):
gdd.download_file_from_google_drive(
file_id='1rZZJRpqQKkncn30Igtk8KirgR96QlCFO',
dest_path=model_path)
chainer.serializers.load_npz(model_path, model)
if args.gpu >= 0:
cuda.get_device_from_id(args.gpu).use()
detector.to_gpu()
model.to_gpu()
chainer.config.train = False
img = utils.read_image(args.image)
# detect persons
bboxes, labels, scores = detector.predict([img])
bbox, label, score = bboxes[0], labels[0], scores[0]
# expand bboxes and crop the image
img = img / 255.
img = img.astype(np.float32)
img_persons = list()
bbox_persons = list()
for ymin, xmin, ymax, xmax in bbox:
scale = ymax - ymin
# this is for ankle (also used in training with mpii dataset)
offset = 15 / 200 * scale
center = (xmin + xmax) / 2, (ymin + ymax) / 2 + offset
# this is for ankle (also used in training with mpii dataset)
scale *= 1.25
xmin, xmax = center[0] - scale / 2, center[0] + scale / 2
ymin, ymax = center[1] - scale / 2, center[1] + scale / 2
# truncate
xmin = int(max(0, xmin))
ymin = int(max(0, ymin))
xmax = int(min(img.shape[2], xmax))
ymax = int(min(img.shape[1], ymax))
# croping
img_person = img[:, ymin:ymax, xmin:xmax]
img_person = transforms.resize(img_person, (256, 256))
img_persons.append(img_person)
bbox_persons.append((ymin, xmin, ymax, xmax))
img_persons = np.array(img_persons)
bbox_persons = np.array(bbox_persons)
utils.write_image(
utils.tile_images((255 * img_persons).astype(np.float32), n_col=2),
'tiled.jpg')
# estimate poses
if args.gpu >= 0:
img_persons = cuda.to_gpu(img_persons)
with chainer.no_backprop_mode():
# (R, 3, 256, 256) -> (R, 16, 64, 64) -> (16, 64, 64)
_outputs, outputs = model(img_persons)
outputs = cuda.to_cpu(outputs.array)
R, C, H, W = outputs.shape
# heatmap to keypoint
# R, C, H, W -> R, C, 2
keypoints = list()
for output in outputs:
# (16, 64, 64) -> (16, )
output = output.reshape(C, -1).argmax(axis=1)
keypoint = np.unravel_index(output, (H, W))
keypoint = np.array(keypoint).T
keypoints.append(keypoint)
# keypoint (local) to keypoint (global)
keypoint_persons = list()
for keypoint, bbox_person in zip(keypoints, bbox_persons):
ymin, xmin, ymax, xmax = bbox_person
keypoint = transforms.resize_point(keypoint, (H, W),
(ymax - ymin, xmax - xmin))
keypoint_person = keypoint + np.array((ymin, xmin))
keypoint_persons.append(keypoint_person)
# visualize
img = cv2.imread(args.image)
visualizer = MPIIVisualizer()
img_pose = img.copy()
for keypoint_person, bbox_person in zip(keypoint_persons, bbox_persons):
ymin, xmin, ymax, xmax = bbox_person
img_pose = visualizer.run(img_pose, keypoint_person)
img_pose = cv2.rectangle(img_pose, (xmin, ymin), (xmax, ymax),
(0, 255, 255), 10)
cv2.imwrite('input.jpg', img)
cv2.imwrite('output.jpg', img_pose)
if args.gpu >= 0:
imgs = xp.asnumpy(imgs.array)
out = xp.asnumpy(out_v.array)
else:
imgs = imgs.array
out = out_v.array
## output images
for i in range(len(out)):
fn = dataset.get_img_path(cnt)
bfn, ext = os.path.splitext(fn)
print("\nProcessing {}".format(fn))
if args.class_num > 0:
write_image((255 * np.stack(
[out[i, 2], np.zeros_like(out[i, 0]), out[i, 1]],
axis=0)).astype(np.uint8),
os.path.join(outdir, os.path.basename(bfn)) +
".jpg")
new = np.argmax(out[i], axis=0)
# print(new.shape)
else:
new = dataset.var2img(out[i])
print("raw value: {} -- {}".format(np.min(out[i]), np.max(out[i])))
print("image value: {} -- {}, ".format(np.min(new), np.max(new),
new.shape))
# converted image
if args.imgtype == "dcm":
path = os.path.join(outdir, os.path.basename(fn))
ref_dicom = dataset.overwrite_dicom(new, fn, salt)
ref_dicom.save_as(path)
elif args.imgtype == "npy":
