Initial operators working with masking

.gitignore

@@ -160,3 +160,5 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+*.DS_Store

README.md

@@ -1,2 +1,10 @@
 # cdim
 Constrained Diffusion Implicit Models
+
+conda create -n cdim python=3.11
+
+conda activate cdim
+
+pip install -r requirements.txt
+
+pip install torch==2.4.1+cu124 torchvision-0.19.1+cu124 --extra-index-url https://download.pytorch.org/whl/cu124

cdim/image_utils.py

@@ -0,0 +1,17 @@
+from torchvision.transforms import ToPILImage
+
+def save_to_image(tensor, filename):
+    """
+    Saves a torch tensor to an image.
+    The image assumed to be (1, 3, H, W)
+    with values between (-1, 1)
+    """
+    to_save = (tensor[0] + 1) / 2
+    to_save = to_save.clamp(0, 1)
+
+    # Convert to PIL Image
+    transform = ToPILImage()
+    img = transform(to_save)
+
+    # Save the image
+    img.save(filename)

cdim/noise.py

@@ -0,0 +1,58 @@
+# Code based on https://github.com/DPS2022/diffusion-posterior-sampling
+from abc import ABC, abstractmethod
+import torch
+
+__NOISE__ = {}
+
+
+def register_noise(name: str):
+    def wrapper(cls):
+        if __NOISE__.get(name, None):
+            raise NameError(f"Name {name} is already defined!")
+        __NOISE__[name] = cls
+        return cls
+    return wrapper
+
+
+def get_noise(name: str, **kwargs):
+    if __NOISE__.get(name, None) is None:
+        raise NameError(f"Name {name} is not defined.")
+    noiser = __NOISE__[name](**kwargs)
+    noiser.__name__ = name
+    return noiser
+
+
+class Noise(ABC):
+    def __call__(self, data):
+        return self.forward(data)
+    
+    @abstractmethod
+    def __call__(self, data):
+        pass
+
+@register_noise(name='gaussian')
+class GaussianNoise(Noise):
+    def __init__(self, sigma):
+        self.sigma = sigma
+    
+    def __call__(self, data):
+        # Important! We scale sigma by 2 because the config assumes images are in [0, 1]
+        # but actually this model uses images in [-1, 1]
+        return data + torch.randn_like(data, device=data.device) * self.sigma * 2
+
+
+@register_noise(name='poisson')
+class PoissonNoise(Noise):
+    def __init__(self, rate):
+        self.rate = rate
+
+    def __call__(self, data):
+        import numpy as np
+        data = (data + 1.0) / 2.0
+        data = data.clamp(0, 1)
+        device = data.device
+        data = data.detach().cpu()
+        data = torch.from_numpy(np.random.poisson(data * 255.0 * self.rate) / 255.0 / self.rate)
+        data = data * 2.0 - 1.0
+        data = data.clamp(-1, 1)
+        return data.to(device)

cdim/noise_configs/gaussian_noise.yaml

@@ -1,2 +0,0 @@
-name: gaussian
-sigma: 0.05

cdim/operators/__init__.py

@@ -0,0 +1,24 @@
+# Code based on https://github.com/DPS2022/diffusion-posterior-sampling
+from abc import ABC, abstractmethod
+
+__OPERATOR__ = {}
+
+
+def register_operator(name: str):
+    def wrapper(cls):
+        if __OPERATOR__.get(name, None):
+            raise NameError(f"Name {name} is already registered!")
+        __OPERATOR__[name] = cls
+        return cls
+    return wrapper
+
+
+def get_operator(name: str, **kwargs):
+    if __OPERATOR__.get(name, None) is None:
+        raise NameError(f"Name {name} is not defined.")
+    return __OPERATOR__[name](**kwargs)
+
+
+# Import everything to make sure they register
+from .random_box_masker import RandomBoxMasker
+from .identity_operator import IdentityOperator

cdim/operators/identity_operator.py

@@ -0,0 +1,9 @@
+from cdim.operators import register_operator
+
+@register_operator(name='identity')
+class IdentityOperator:
+    def __init__(self, device):
+        self.device = device
+    
+    def __call__(self, data):
+        return data

cdim/operators/random_box_masker.py

@@ -0,0 +1,56 @@
+import torch
+
+from cdim.operators import register_operator
+
+@register_operator(name='box_inpainting')
+class RandomBoxMasker:
+    def __init__(self, height=256, width=256, channels=3, box_size=128, device='cpu'):
+        """
+        Initialize the ConsistentBoxMasker with random box positioning.
+        
+        Args:
+        height (int): Height of the input tensors (default: 256)
+        width (int): Width of the input tensors (default: 256)
+        channels (int): Number of channels in the input tensors (default: 3)
+        box_size (int): Size of the box to mask (default: 128)
+        device (str): Device to create the mask on (default: 'cpu')
+        """
+        self.height = height
+        self.width = width
+        self.channels = channels
+        self.box_size = min(box_size, height, width)  # Ensure box_size doesn't exceed image dimensions
+        self.device = device
+        
+        # Create a binary mask for box selection
+        self.mask = torch.ones((1, channels, height, width), device=device)
+        
+        # Randomly calculate the top-left corner of the box
+        max_y = height - self.box_size
+        max_x = width - self.box_size
+
+        start_y = torch.randint(0, max_y + 1, (1,)).item()
+        start_x = torch.randint(0, max_x + 1, (1,)).item()
+        
+        # Set the box area in the mask to 0
+        self.mask[0, :, start_y:start_y+self.box_size, start_x:start_x+self.box_size] = 0
+    
+    def __call__(self, tensor):
+        """
+        Apply the consistent box masking to the input tensor.
+        
+        Args:
+        tensor (torch.Tensor): Input tensor of shape (b, channels, height, width)
+        
+        Returns:
+        torch.Tensor: Tensor with the same shape as input, but with the box area masked out
+        """
+        b, c, h, w = tensor.shape
+        assert c == self.channels and h == self.height and w == self.width, \
+            f"Input tensor must be of shape (b, {self.channels}, {self.height}, {self.width})"
+        
+        # Move the mask to the same device as the input tensor if necessary
+        if tensor.device != self.mask.device:
+            self.mask = self.mask.to(tensor.device)
+        
+        # Apply the mask to the input tensor
+        return tensor * self.mask

inference.py

@@ -1,12 +1,55 @@
 import argparse
 import os
+import yaml
+
+from PIL import Image
+import numpy as np
+import torch
+
+from cdim.noise import get_noise
+from cdim.operators import get_operator
+from cdim.image_utils import save_to_image
+
+
+def load_image(path):
+    """
+    Load the image and normalize to [-1, 1]
+    """
+    original_image = Image.open(path)
+
+    # Resize if needed
+    original_image = np.array(original_image.resize((256, 256), Image.BICUBIC))
+    original_image = torch.from_numpy(original_image).unsqueeze(0).permute(0, 3, 1, 2)
+    return (original_image / 127.5 - 1.0).to(torch.float)
+    
+
+def load_yaml(file_path: str) -> dict:
+    with open(file_path) as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+    return config
 
 
 def main(args):
+    device_str = f"cuda" if args.cuda and torch.cuda.is_available() else 'cpu'
+    print(f"Using device {device_str}")
+    device = torch.device(device_str) 
+
     os.makedirs(args.output_dir, exist_ok=True)
+    original_image = load_image(args.input_image).to(device)
 
-    
+    # Load the noise function
+    noise_config = load_yaml(args.noise_config)
+    noise_function = get_noise(**noise_config)
+    print(noise_function)    
 
+    # Load the measurement function A
+    operator_config = load_yaml(args.operator_config)
+    operator_config["device"] = device
+    operator = get_operator(**operator_config)
+    print(operator)
+
+    noisy_measurement = noise_function(operator(original_image))
+    save_to_image(noisy_measurement, os.path.join(args.output_dir, "noisy_measurement.png"))
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
@@ -17,4 +60,6 @@ if __name__ == '__main__':
     parser.add_argument("operator_config", type=str)
     parser.add_argument("noise_config", type=str)
     parser.add_argument("--output-dir", default=".", type=str)
+    parser.add_argument("--cuda", default=True, action=argparse.BooleanOptionalAction)
+
     main(parser.parse_args())

noise_configs/gaussian_noise_config.yaml

@@ -0,0 +1,3 @@
+name: gaussian
+sigma: 0.05
+# Important! This noise is assumed to be for images in [0, 1]

operator_configs/box_inpainting_config.yaml

@@ -0,0 +1,5 @@
+name: box_inpainting
+box_size: 128
+height: 256
+width: 256
+channels: 3

operator_configs/identity_operator_config.yaml

@@ -0,0 +1 @@
+name: identity

requirements.txt

@@ -0,0 +1,2 @@
+numpy==2.1.2
+Pillow==11.0.0