File size: 2,097 Bytes
78a947a |
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 |
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import torch
from src.datasets.utils import dequantize_16_bit
from src.models.compound_model import CompoundModel
from src.models.conv3don2d import Conv3Don2D
from src.models.ensemble_model import EnsembleModel
from src.models.unet import UNet
base_bath = Path("./model_weights")
input_channels = 124 // 2
out_channels = 1
conv3d_filters = 1
conv3d_kernel_size = (7, 2, 2)
unet_input_channels = (input_channels - (conv3d_kernel_size[0] - 1)) * conv3d_filters
models = []
for i in range(1, 6):
compression_model = Conv3Don2D(
in_channels=input_channels,
num_filters=conv3d_filters,
kernel_size=conv3d_kernel_size,
normalization=None,
)
unet = UNet(
in_channels=unet_input_channels,
out_channels=out_channels,
pad=False,
bilinear=True,
normalization=None,
)
m = CompoundModel(compression_model, unet)
weights = torch.load(base_bath / f"{i}" / f"best_weights.pt", weights_only=True)
m.load_state_dict(weights, strict=False)
models.append(m)
# Create an ensemble model with the loaded models
ensemble_model = EnsembleModel(*models)
"""
Example usage:
# Predict a sample image using the ensemble model.
# Placeholder for the path to the data
data_path = Path("")
sample = "116-1"
img = np.load(data_path / "images" / f"{sample}.npy")
quant_bias = np.load(data_path / "quant_biases" / f"{sample}.npy")
quant_scale = np.load(data_path / "quant_scales" / f"{sample}.npy")
img = dequantize_16_bit(img, quant_bias, quant_scale)
img = torch.from_numpy(img)
mask = np.load(data_path / "masks" / f"{sample}.npy").squeeze()
# Expects the image to already be padded or cropped to the expected input size.
pred = ensemble_model(img)
pred = pred.detach().cpu().numpy().squeeze()
# Plot the prediction and the masked prediction
plt.imshow(pred)
plt.colorbar()
plt.savefig("ensemble_prediction.png")
plt.clf()
masked_pred = pred * mask
plt.imshow(masked_pred)
plt.colorbar()
plt.savefig("ensemble_prediction_masked.png")
""" |