def train_transform(rgb, depth):
s = np.random.uniform(1.0, 1.5) # random scaling
# print("scale factor s={}".format(s))
depth_np = depth / s
angle = np.random.uniform(-5.0, 5.0) # random rotation degrees
do_flip = np.random.uniform(0.0, 1.0) < 0.5 # random horizontal flip
# perform 1st part of data augmentation
transform = transforms.Compose([
transforms.Crop(130, 10, 240, 1200),
transforms.Rotate(angle),
transforms.Resize(s),
transforms.CenterCrop((oheight, owidth)),
transforms.HorizontalFlip(do_flip)
])
rgb_np = transform(rgb)
# random color jittering
rgb_np = color_jitter(rgb_np)
rgb_np = np.asfarray(rgb_np, dtype='float') / 255
# Scipy affine_transform produced RuntimeError when the depth map was
# given as a 'numpy.ndarray'
depth_np = np.asfarray(depth_np, dtype='float32')
depth_np = transform(depth_np)
return rgb_np, depth_np
def train_transform(rgb, depth):
s = np.random.uniform(1.0, 1.5) # random scaling
# print("scale factor s={}".format(s))
depth_np = depth / s
angle = np.random.uniform(-5.0, 5.0) # random rotation degrees
do_flip = np.random.uniform(0.0, 1.0) < 0.5 # random horizontal flip
# perform 1st part of data augmentation
transform = transforms.Compose([
transforms.Resize(
250.0 / iheight
), # this is for computational efficiency, since rotation is very slow
transforms.Rotate(angle),
transforms.Resize(s),
transforms.CenterCrop((oheight, owidth)),
transforms.HorizontalFlip(do_flip)
])
rgb_np = transform(rgb)
# random color jittering
rgb_np = color_jitter(rgb_np)
rgb_np = np.asfarray(rgb_np, dtype='float') / 255
depth_np = transform(depth_np)
return rgb_np, depth_np
def train_transform(rgb, depth):
s = np.random.uniform(1.0, 1.5) # random scaling
# print("scale factor s={}".format(s))
depth_np = depth / s
angle = np.random.uniform(-5.0, 5.0) # random rotation degrees
do_flip = np.random.uniform(0.0, 1.0) < 0.5 # random horizontal flip
# set zeros in depth as NaN
depth_np[depth_np == 0] = np.nan
# perform 1st part of data augmentation
transform = transforms.Compose([
transforms.Resize(
float(image_size) / iheight
), # this is for computational efficiency, since rotation is very slow
transforms.Rotate(angle),
transforms.Resize(s),
transforms.CenterCrop((oheight, owidth)),
transforms.HorizontalFlip(do_flip),
])
rgb_np = transform(rgb)
# random color jittering
rgb_np = color_jitter(rgb_np)
rgb_np = np.asfarray(rgb_np, dtype='float') / 255
rgb_np = normalize(rgb_np) # from [0,1] to [-1,1]
depth_np = transform(depth_np)
depth_np[np.isnan(depth_np)] = 0
depth_np = depth_np / 10.0
return rgb_np, depth_np
def train_transform(self, rgb: np.ndarray, depth_raw: np.ndarray,
depth_fix: np.ndarray) -> TNpData:
s = np.random.uniform(1.0, 1.5) # random scaling
depth_raw = depth_raw / s
depth_fix = depth_fix / s
angle = np.random.uniform(-5.0, 5.0) # random rotation degrees
do_flip = np.random.uniform(0.0, 1.0) < 0.5 # random horizontal flip
# perform 1st part of data augmentation
transform = transforms.Compose([
transforms.Resize(
250.0 / self.iheight
), # this is for computational efficiency, since rotation is very slow
transforms.Rotate(angle),
transforms.Resize(s),
transforms.CenterCrop((self.oheight, self.owidth)),
transforms.HorizontalFlip(do_flip)
])
rgb = transform(rgb)
# random color jittering
rgb = color_jitter(rgb)
rgb = np.asfarray(rgb, dtype='float') / 255
depth_raw = transform(depth_raw)
depth_fix = transform(depth_fix)
return rgb, depth_raw, depth_fix
def train_transform(rgb, depth):
s = np.random.uniform(1.0, 1.5) # random scaling
# print("scale factor s={}".format(s))
depth_np = depth / s
angle = np.random.uniform(-5.0, 5.0) # random rotation degrees
do_flip = np.random.uniform(0.0, 1.0) < 0.5 # random horizontal flip
# perform 1st part of data augmentation
transform = transforms.Compose([
#transforms.Resize(530 / iheight), # this is for computational efficiency, since rotation is very slow
transforms.Resize(250 / iheight),
transforms.Rotate(angle),
transforms.Resize(s),
transforms.CenterCrop((oheight, owidth)),
transforms.HorizontalFlip(do_flip)
])
rgb_np = transform(rgb)
# 自己添加
# rgb_np = Transform.resize(rgb_np, [512, 512])
rgb_np = cv2.resize(rgb_np, (512, 512), interpolation=cv2.INTER_NEAREST)
###########
# random color jittering
rgb_np = color_jitter(rgb_np)
rgb_np = np.asfarray(rgb_np, dtype='float') / 255
depth_np = transform(depth_np)
#自己添加
depth_np = cv2.resize(depth_np, (512, 512),
interpolation=cv2.INTER_NEAREST)
#depth_np=Transform.resize(depth_np,[512,512])
###########
#data=rgb_np*255
#data=Image.fromarray(data.astype(np.uint8))
#data.show()
return rgb_np, depth_np
def main():
in_z = 0
volpath = os.path.join(args.datapath, 'train')
segpath = volpath
batch_size = args.batch_size
orig_dim = 256
sqr = transforms.Square()
aff = transforms.Affine()
crop = transforms.RandomCrop(224)
scale = transforms.Scale(orig_dim)
rotate = transforms.Rotate(0.5, 30)
noise = transforms.Noise(0.02)
flip = transforms.Flip()
transform_plan = [sqr, scale, aff, rotate, crop, flip, noise]
lr = 1e-4
series_names = ['Mag']
seg_series_names = ['AV']
model = models.Net23(2)
model.cuda()
optimizer = optim.RMSprop(model.parameters(), lr=lr)
out_z, center_crop_sz = utils.get_out_size(orig_dim, in_z,\
transform_plan, model)
t0 = time.time()
counter = 0
print_interval = args.log_interval
model.train()
for i in range(200000000000000000000000000000000):
weight = torch.FloatTensor([0.2, 0.8]).cuda()
vol, seg, inds = preprocess.get_batch(volpath, segpath, batch_size, in_z,\
out_z, center_crop_sz, series_names, seg_series_names,\
transform_plan, 8, nrrd=True)
vol = torch.unsqueeze(vol, 1)
vol = Variable(vol).cuda()
seg = Variable(seg).cuda()
out = model(vol).squeeze()
loss = F.cross_entropy(out, seg, weight=weight)
optimizer.zero_grad()
loss.backward()
optimizer.step()
counter += 1
sys.stdout.write('\r{:.2f}%'.format(counter * batch_size /
print_interval))
sys.stdout.flush()
if counter * batch_size >= print_interval and i > 0:
seg_hot = utils.get_hot(seg.data.cpu().numpy(), out.size()[-3])
seg_hot = np.transpose(seg_hot, axes=[1, 0, 2, 3])
out_hot = np.argmax(out.data.cpu().numpy(), axis=1)
out_hot = utils.get_hot(out_hot, out.size()[-3])
out_hot = np.transpose(out_hot, axes=[1, 0, 2, 3])
dce2 = utils.dice(seg_hot[:, 1:], out_hot[:, 1:])
vol_ind = utils.get_liver(seg)
v = vol.data.cpu().numpy()[vol_ind].squeeze()
real_seg = seg[vol_ind].data.cpu().numpy()
out_plt = out.data.cpu().numpy()[vol_ind]
out_plt = np.argmax(out_plt, axis=0)
fake_seg = out_plt
masked_real = utils.makeMask(v, real_seg, out.size()[-3], 0.5)
masked_fake = utils.makeMask(v, fake_seg, out.size()[-3], 0.5)
fig = plt.figure(1)
fig.suptitle('Volume {} ; Dice = {:.2f}'.format(\
inds[vol_ind],dce2))
v = fig.add_subplot(1, 2, 1)
v.set_title('real')
plt.imshow(masked_real)
sreal = fig.add_subplot(1, 2, 2)
sreal.set_title('fake')
plt.imshow(masked_fake)
outfile = os.path.join(args.datapath, 'out.png')
plt.savefig(outfile, dpi=200)
plt.clf()
print(('\rIteration {}: Block completed in {:.2f} sec ; Loss = {:.2f}')\
.format(i*batch_size,time.time()-t0, loss.data.cpu()[0]))
checkpoint_file = os.path.join(args.datapath,\
'model_checkpoint.pth')
torch.save(model.state_dict(), checkpoint_file)
counter = 0
t0 = time.time()
import asf_read as reader
import numpy as np
# take a dataset, augment it, and write it back into an asf
train_dataset = FaceLandmarksDataset(root_dir='imm_face_db/training/')
test_dataset = FaceLandmarksDataset(root_dir='imm_face_db/testing/')
batch_size = 8
train_loader = DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True)
test_loader = DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=False)
available_transformations = {
'rotate': tr.Rotate(),
'jitter': tr.Jitter(),
'horizontal_shift': tr.HorizontalShift(),
'vertical_shift': tr.VerticalShift()
}
def generate_transforms():
num_to_apply = random.randint(0, 2)
num_applied = 0
transformation_list = []
# create a list of transformations
while (num_applied < num_to_apply):
key = random.choice(list(available_transformations))
transformation = available_transformations[key]
checkpoint_dir = 'checkpont directory'
volpath = 'path to image volumes'
segpath = 'path to segmentation masks'
num_labels = 3 #labels: lateral annulus, medial annulus, background
crop_size = 224
original_size = 256
num_steps = 1000000000 #number of training steps
################################################################
sqr = transforms.Square()
aff = transforms.Affine()
crop = transforms.RandomCrop(crop_size)
scale = transforms.Scale(original_size)
rotate = transforms.Rotate(0.5, 30)
noise = transforms.Noise(0.02)
transform_plan = [sqr, scale, aff, rotate, crop, noise]
lr = 1e-4
series_names = ['echo']
seg_series_names = ['echo']
model = models.Net23(3)
model.cuda()
optimizer = optim.RMSprop(model.parameters(), lr=lr)
out_z, center_crop_sz = utils.get_out_size(original_size, 0,\
transform_plan, model)