if __name__ == "__main__":
image_shape = (255, 255)
data_root = './data/ILSVRC/ILSVRC2015'
transform = transforms.Compose([transforms.ToTensor()])
test_set = ILSVRC(data_root,
image_shape=image_shape,
data_percentage=0.5,
train=False,
transform=transform)
model = ConvVAE()
model.load_state_dict(torch.load("./trained_models/ConvVAE.pt"))
model.eval()
model2 = ConvVAE()
model2.load_state_dict(
torch.load("./trained_models/ConvVAE_firstConv_GMN_batch.pt"))
model2.eval()
index = random.randint(0, len(test_set))
print(index)
images, _, ground_truth, count, resized_template = test_set[index]
templates = torch.reshape(
resized_template,
(1, resized_template.shape[-3], resized_template.shape[-2],
resized_template.shape[-1]))
decoded, _, _ = model(templates)
im1 = transforms.ToPILImage()(templates[0]).convert("RGB")
transforms.ToTensor()
])
test_set = datasets.CIFAR10(root=dataset_root,
train=False,
download=True,
transform=transform)
model_r = ConvVAE(channels=1)
model_g = ConvVAE(channels=1)
model_b = ConvVAE(channels=1)
model_r.load_state_dict(torch.load("./trained_models/ConvVAE_r.pt"))
model_g.load_state_dict(torch.load("./trained_models/ConvVAE_g.pt"))
model_b.load_state_dict(torch.load("./trained_models/ConvVAE_b.pt"))
model_r.eval()
model_g.eval()
model_b.eval()
index = random.randint(0, len(test_set))
print(index)
image, class_index = test_set[4942]
image = torch.reshape(
image, (1, image.shape[-3], image.shape[-2], image.shape[-1]))
decoded_r, _, _ = model_r(image[:, 0, :, :].unsqueeze_(0))
decoded_g, _, _ = model_g(image[:, 1, :, :].unsqueeze_(0))
decoded_b, _, _ = model_b(image[:, 2, :, :].unsqueeze_(0))
decoded = torch.cat((decoded_r, decoded_g, decoded_b), 1)
im1 = transforms.ToPILImage()(image[0]).convert("RGB")