bcmi / duconet-image-harmonization Goto Github PK

        transform = transforms.Compose([
            transforms.ToTensor(),
         ])
        print(type(image_lab))
        comp_image_lab = transform(image_lab).unsqueeze(0).to(self.device)
        predicted_output = self.net(image, mask,image_lab = comp_image_lab)
        predicted_image = predicted_output['images']


        for transform in reversed(self.transforms):
            predicted_image = transform.inv_transform(predicted_image)

        predicted_image = torch.clamp(predicted_image, 0, 255)

    if return_numpy:
        return predicted_image.cpu().numpy()
    else:
        return predicted_image