Пример #1
0
d1_t, d2_t = data.get_dataset_triangles("test",
                                        data_size=50,
                                        hollow=False,
                                        batch_size=batch_size)

lmbda = 2048
random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)
np.random.seed(1)
print("=" * 50)
net = inverseConsistentNet.InverseConsistentNet(
    network_wrappers.DoubleNet(
        network_wrappers.RandomShift(0.25),
        network_wrappers.FunctionFromVectorField(
            networks.tallUNet2(dimension=2)),
    ),
    lambda x, y: torch.mean((x - y)**2),
    lmbda,
)

input_shape = next(iter(d1))[0].size()
network_wrappers.assignIdentityMap(net, input_shape)
net.cuda()
optimizer = torch.optim.Adam(net.parameters(), lr=0.001)
net.train()

xs = []
for _ in range(40):
    y = np.array(train.train2d(net, optimizer, d1, d2, epochs=50))
    xs.append(y)
Пример #2
0
import networks
import network_wrappers
import data
import describe

BATCH_SIZE = 32
SCALE = 1  # 1 IS QUARTER RES, 2 IS HALF RES, 4 IS FULL RES
input_shape = [BATCH_SIZE, 1, 40 * SCALE, 96 * SCALE, 96 * SCALE]

GPUS = 4


phi = network_wrappers.FunctionFromVectorField(
    networks.tallUNet(unet=networks.UNet2ChunkyMiddle, dimension=3)
)
psi = network_wrappers.FunctionFromVectorField(networks.tallUNet2(dimension=3))

pretrained_lowres_net = inverseConsistentNet.InverseConsistentNet(
    network_wrappers.DoubleNet(phi, psi),
    lambda x, y: torch.mean((x - y) ** 2),
    100,
)

network_wrappers.assignIdentityMap(pretrained_lowres_net, input_shape)


network_wrappers.adjust_batch_size(pretrained_lowres_net, 12)
trained_weights = torch.load(
    "results/dd_l400_continue_rescalegrad2/knee_aligner_resi_net1800"
)
from mermaidlite import compute_warped_image_multiNC, identity_map_multiN
import torch
import random
import inverseConsistentNet
import networks
import data
import describe

BATCH_SIZE = 8
SCALE = 2  # 1 IS QUARTER RES, 2 IS HALF RES, 4 IS FULL RES
working_shape = [BATCH_SIZE, 1, 40 * SCALE, 96 * SCALE, 96 * SCALE]

GPUS = 4

net = inverseConsistentNet.InverseConsistentNet(
    networks.tallUNet2(dimension=3),
    lmbda=512,
    input_shape=working_shape,
    random_sampling=False,
)

knees, medknees = data.get_knees_dataset()
knees = [F.avg_pool3d(knee, 2) for knee in knees]
if GPUS == 1:
    net_par = net.cuda()
else:
    net_par = torch.nn.DataParallel(net).cuda()
optimizer = torch.optim.Adam(net_par.parameters(), lr=0.00005)

net_par.train()
GPUS = 4

tmp_affine_net = inverseConsistentNet.InverseConsistentAffineNet(
    networks.ConvolutionalMatrixNet(dimension=3),
    lmbda=100,
    input_shape=working_shape,
)

phi = networks.StumpyConvolutionalMatrixNet(dimension=3)
psi = networks.StumpyConvolutionalMatrixNet(dimension=3)
affine_net = networks.DownscaleConvolutionalMatrixNet(
    networks.DoubleAffineNet(phi, psi, tmp_affine_net.identityMap,
                             tmp_affine_net.spacing))

deformable_phi = networks.DownsampleNet(networks.tallUNet2(dimension=3),
                                        dimension=3)
deformable_psi = networks.tallUnet2(dimension=3)
deformable_net = networks.DoubleDeformableNet(deformable_phi, deformable_psi,
                                              tmp)
net = inverseConsistentNet.InverseConsistentAffineDeformableNet(
    affine_net,
    networks.tallUNet2(dimension=3),
    lmbda=100,
    input_shape=working_shape)

pretrained_weights = torch.load(
    "results/affine_knee_double_stumpy10/knee_aligner_resi_net35400")
pretrained_weights = OrderedDict([
    (a.split("regis_net.")[1], b) for a, b in pretrained_weights.items()
    if ("regis_net" in a and not ("identityMap$" in a + "$"))
import random
import os

import describe

d1_triangles, d2_triangles = data.get_dataset_triangles("train",
                                                        data_size=50,
                                                        hollow=True,
                                                        samples=128)
d1_triangles_test, d2_triangles_test = data.get_dataset_triangles("test",
                                                                  data_size=50,
                                                                  hollow=True,
                                                                  samples=256)

# d1_triangles_test, d2_triangles_test =d1_triangles, d2_triangles
network1 = networks.tallUNet2()
network2 = networks.tallUNet2()

network = networks.TwoStepNet

d1, d2, d1_t, d2_t = (d1_triangles, d2_triangles, d1_triangles_test,
                      d2_triangles_test)
lmbda = 2048
random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)
np.random.seed(1)
print("=" * 50)
print(network, lmbda)
net = inverseConsistentNet.InverseConsistentNet(network(), lmbda,
                                                next(iter(d1))[0].size())
Пример #6
0
                                    data_size=data_size,
                                    hollow=True,
                                    batch_size=batch_size)
d1_t, d2_t = data.get_dataset_triangles("test",
                                        data_size=data_size,
                                        hollow=True,
                                        batch_size=batch_size)

image_A, image_B = (x[0].cuda() for x in next(zip(d1, d2)))

net_tmp = inverseConsistentNet.InverseConsistentAffineNet(
    networks.ConvolutionalMatrixNet(), 100,
    next(iter(d1))[0].size())

net = inverseConsistentNet.InverseConsistentAffineNet(
    networks.AffineFromUNet(networks.tallUNet2(), net_tmp.identityMap),
    100,
    next(iter(d1))[0].size(),
)

net.cuda()

import train

optim = torch.optim.Adam(net.parameters(), lr=0.00001)
net.train().cuda()

xs = []
for _ in range(240):
    y = np.array(train.train2d(net, optim, d1, d2, epochs=50))
    xs.append(y)