import footsteps

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_net25500"
)

# trained_weights = torch.load("../results/dd_knee_l400_continue_smallbatch2/knee_aligner_resi_net9300")
# rained_weights = torch.load("../results/double_deformable_knee3/knee_aligner_resi_net22200")
pretrained_lowres_net.load_state_dict(trained_weights)
data_size = 50
d1, d2 = data.get_dataset_triangles("train",
                                    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 = inverseConsistentNet.InverseConsistentNet(
    network_wrappers.DoubleNet(
        network_wrappers.FunctionFromVectorField(
            networks.tallUNet2(dimension=2)),
        network_wrappers.FunctionFromVectorField(
            networks.tallUNet2(dimension=2)),
    ),
    lambda x, y: torch.mean((x - y)**2),
    700,
)

input_shape = next(iter(d1))[0].size()
network_wrappers.assignIdentityMap(net, input_shape)
net.cuda()

import icon_registration.train as train

optim = torch.optim.Adam(net.parameters(), lr=0.0001)
net.train().cuda()
import icon_registration.network_wrappers as network_wrappers
import icon_registration.data as data
import footsteps

BATCH_SIZE = 4
SCALE = 2  # 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))
net = inverseConsistentNet.InverseConsistentNet(
    network_wrappers.DoubleNet(
        network_wrappers.DownsampleNet(network_wrappers.DoubleNet(phi, psi),
                                       dimension=3),
        network_wrappers.FunctionFromVectorField(
            networks.tallUNet2(dimension=3)),
    ),
    inverseConsistentNet.ncc,
    165000,
)

network_wrappers.assignIdentityMap(net, input_shape)

weights = torch.load("results/hires_ncc_70000_6/knee_aligner_resi_net4200")
net.load_state_dict(weights)

if GPUS == 1:
    net_par = net.cuda()
else:
    net_par = torch.nn.DataParallel(net).cuda()
示例#4
0
import icon_registration.data as data
import footsteps

BATCH_SIZE = 12
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 = network_wrappers.DoubleNet(phi, psi)

hires_net = inverseConsistentNet.InverseConsistentNet(
    network_wrappers.DoubleNet(
        network_wrappers.DownsampleNet(pretrained_lowres_net, dimension=3),
        network_wrappers.FunctionFromVectorField(networks.tallUNet2(dimension=3)),
    ),
    inverseConsistentNet.ssd_only_interpolated,
    3600,
)
BATCH_SIZE = 4
SCALE = 2  # 1 IS QUARTER RES, 2 IS HALF RES, 4 IS FULL RES
input_shape = [BATCH_SIZE, 1, 40 * SCALE, 96 * SCALE, 96 * SCALE]
network_wrappers.assignIdentityMap(hires_net, input_shape)

trained_weights = torch.load("results/hires_smart_6/knee_aligner_resi_net74700")
data_size = 50
d1, d2 = data.get_dataset_triangles("train",
                                    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 = inverseConsistentNet.InverseConsistentNet(
    network_wrappers.DoubleNet(
        network_wrappers.FunctionFromMatrix(
            networks.ConvolutionalMatrixNet(dimension=2), ),
        network_wrappers.FunctionFromVectorField(
            networks.tallUNet2(dimension=2)),
    ),
    lambda x, y: torch.mean((x - y)**2),
    700,
)

input_shape = next(iter(d1))[0].size()
network_wrappers.assignIdentityMap(net, input_shape)
# pretrained_weights = torch.load("results/affine_triangle_pretrain/epoch000case0.png")
# pretrained_weights = OrderedDict(
#    [
#        (a.split("regis_net.")[1], b)
#        for a, b in pretrained_weights.items()
#        if "regis_net" in a
#    ]