data and checkpoint added

.gitignore

@@ -0,0 +1,164 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+#jupyter
+.ipynb_checkpoints
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  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/

__init__.py

@@ -0,0 +1,23 @@
+from bikefusion import *
+from torch.utils.data import TensorDataset, random_split
+import torch
+import os
+
+current_dir = os.path.dirname(os.path.abspath(__file__))
+
+def load_bikefusion_and_data():
+    partial_images, masks, parametric, description, targets  = load_data(os.path.join(current_dir, 'data/'))
+    
+    training_images = preprocess(targets)
+    
+    dataset = TensorDataset(training_images)
+
+    # split to training and validation
+    train_size = int(0.8 * len(dataset))
+    val_size = len(dataset) - train_size
+    train_dataset, val_dataset = random_split(dataset, [train_size, val_size], generator=torch.Generator().manual_seed(42))
+    
+    Diffuser = InpaintingDenoisingDiffusion(train_dataset, val_dataset, image_size=128)
+    Diffuser.load_checkpoint(os.path.join(current_dir, 'chekpoint/bikefusion.pt'))
+    
+    return Diffuser

bikefusion/__init__.py

@@ -0,0 +1,4 @@
+from .data_utils import load_data, preprocess, postprocess, to_mask_map
+from .visualizers import visualize_imagesets, visualize_image_evolution
+from .pipline import InpaintingDenoisingPipeline
+from .diffusion import InpaintingDenoisingDiffusion

bikefusion/data_utils.py

@@ -0,0 +1,99 @@
+import torchvision.transforms.functional as F
+import torch
+import numpy as np
+import pickle
+
+def to_mask_map(maks, image_size=(80,128)):
+    bs = maks.shape[0]
+    mask_map = np.ones((bs, 1, image_size[0], image_size[1]))
+    
+    for i in range(bs):
+        left, bottom, right, top = maks[i]
+        mask_map[i, 0, top:bottom, left:right] = 0
+        
+    return mask_map
+
+def pad_to_square(image, target_size=(128, 128)):
+    # Calculate padding for each dimension
+    h, w = image.shape[1], image.shape[2]
+    delta_w = target_size[1] - w
+    delta_h = target_size[0] - h
+    padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
+    
+    # Pad with white (255) for each side
+    padded_image = F.pad(image, padding, fill=255)
+    return padded_image
+
+def preprocess(images):
+    # Convert arrays to tensors
+    images = torch.from_numpy(images).float()
+
+    # Apply padding to each image in the dataset
+    images = torch.stack([pad_to_square(img) for img in images])
+    
+    # Normalize images to [-1, 1]
+    images = images / 255.0 * 2 - 1
+    
+    return images
+
+def un_pad(image, target_size=(80, 128)):
+    # Calculate padding for each dimension
+    h, w = image.shape[1], image.shape[2]
+    delta_w = w - target_size[1]
+    delta_h = h - target_size[0]
+    
+    # Unpad the image
+    un_padded_image = image[:, delta_h // 2:h - delta_h // 2, delta_w // 2:w - delta_w // 2]
+    
+    return un_padded_image
+
+def postprocess(images):
+    # Convert tensors to arrays
+    images = images.detach().cpu().numpy()
+
+    # Unpad each image in the dataset
+    images = np.stack([un_pad(img) for img in images])
+    
+    # Rescale images to [0, 255]
+    images = (images + 1) / 2 * 255
+    
+    return images.astype(np.uint8)
+
+def load_data(split="train", path="data/"):
+    """
+    A function to load the data.
+
+    Parameters:
+    - split: str, "train" or "test"
+
+    Returns:
+    - masked_images: np.ndarray, shape (n_images, channels, height, width), the images with random masks applied
+    - masks: np.ndarray, shape (n_images, 4), the boundaries of the masks (left, bottom, right, top)
+    - parametric: np.ndarray, shape (n_images, 3), the parametric representation of the images
+    - description: list, the description of the images
+    - images: np.ndarray, shape (n_images, channels, height, width), the original images
+    """
+    masked_images = []
+    if split == "train":
+        for i in range(5):
+            masked_im_slice = np.load(f"{path}masked_train_{i}.npy")
+            masked_images.append(masked_im_slice)
+        masked_images = np.concatenate(masked_images, axis=0)
+
+        images = []
+        for i in range(5):
+            im_slice = np.load(f"{path}images_train_{i}.npy")
+            images.append(im_slice)
+        images = np.concatenate(images, axis=0)
+    else:
+        masked_images = np.load(f"{path}masked_test.npy")
+        images = None
+
+    description = pickle.load(open(f"{path}desc_{split}.pkl", "rb"))
+
+    parametric = np.load(f"{path}param_{split}.npy")
+
+    masks = np.load(f"{path}mask_{split}.npy")
+
+
+    return masked_images, masks, parametric, description, images

bikefusion/diffusion.py

@@ -0,0 +1,254 @@
+from diffusers import UNet2DModel
+from diffusers import DDPMScheduler
+from accelerate import Accelerator
+from diffusers.optimization import get_cosine_schedule_with_warmup
+from tqdm.auto import tqdm
+import torch.nn.functional as F
+import torch
+from torch.utils.data import DataLoader, random_split
+from .pipline import InpaintingDenoisingPipeline as ConditionalDenoisingPipeline
+import numpy as np
+import matplotlib.pyplot as plt
+from .visualizers import visualize_imagesets
+
+class InpaintingDenoisingDiffusion:
+    def __init__(self,
+                train_dataset,
+                validation_dataset,
+                image_size = 128,  # the generated image resolution
+                n_train_noise_timesteps = 1000,  # the number of timesteps for the noise scheduler
+                train_batch_size = 16,
+                eval_batch_size = 16,  # how many images to sample during evaluation
+                num_epochs = 50,
+                gradient_accumulation_steps = 1,
+                learning_rate = 5e-5,
+                lr_warmup_steps = 500,
+                mixed_precision = "fp16",  # `no` for float32, `fp16` for automatic mixed precision
+                masking_range = ((24, 104), (0, 128)),  # the range of the random masks applied in training
+                minimum_mask_portion = 0.4,  # the minimum portion of the image to be masked
+                maximum_mask_portion = 0.8,  # the maximum portion of the image to be masked
+                full_image_probability = 0.5,  # the probability of applying a full image mask (no image)
+                device = None,  # "cuda" or "cpu"
+                model = None,
+                ):
+        
+        self.train_dataset = train_dataset
+        self.validation_dataset = validation_dataset
+        self.image_size = image_size
+        self.train_batch_size = train_batch_size
+        self.eval_batch_size = eval_batch_size
+        self.num_epochs = num_epochs
+        self.gradient_accumulation_steps = gradient_accumulation_steps
+        self.learning_rate = learning_rate
+        self.lr_warmup_steps = lr_warmup_steps
+        self.mixed_precision = mixed_precision
+        self.device = device
+        self.masking_range = masking_range
+        self.full_image_probability = full_image_probability
+        self.minimum_mask_portion = minimum_mask_portion
+        self.maximum_mask_portion = maximum_mask_portion
+        self.n_train_noise_timesteps = n_train_noise_timesteps
+        
+        self.dataloader = DataLoader(self.train_dataset, batch_size=self.train_batch_size, shuffle=True)
+        self.val_loader = DataLoader(self.validation_dataset, batch_size=self.eval_batch_size, shuffle=False)
+        
+        if self.device is None:
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        if model is not None:
+            self.unet = model
+        else:
+            self.unet =  UNet2DModel(
+                    sample_size=self.image_size,  # the target image resolution
+                    in_channels=6,  # the number of input channels, 3 for RGB masked images and 3 for RGB noise
+                    out_channels=3,  # the number of output channels (RGB)
+                    layers_per_block=2,  # how many ResNet layers to use per UNet block
+                    block_out_channels=(128, 256, 512, 768),  # the number of output channels for each UNet block
+                    down_block_types=(
+                        "DownBlock2D", # a regular ResNet downsampling block
+                        "AttnDownBlock2D",  # a ResNet downsampling block with spatial self-attention
+                        "AttnDownBlock2D", 
+                        "AttnDownBlock2D"
+                    ),
+                    up_block_types=(
+                        "AttnUpBlock2D",# a ResNet upsampling block with spatial self-attention
+                        "AttnUpBlock2D",
+                        "AttnUpBlock2D",
+                        "UpBlock2D"# a regular ResNet upsampling block
+                    ),
+                ).to(self.device)
+        
+        self.noise_scheduler = DDPMScheduler(num_train_timesteps=n_train_noise_timesteps)
+        self.optimizer = torch.optim.AdamW(self.unet.parameters(), lr=self.learning_rate)
+        self.lr_scheduler = get_cosine_schedule_with_warmup(
+            optimizer=self.optimizer,
+            num_warmup_steps=self.lr_warmup_steps,
+            num_training_steps=(len(self.dataloader) * self.num_epochs),
+        )
+        self.current_epoch = 0
+        
+        self.denoiser = ConditionalDenoisingPipeline(self.unet, self.noise_scheduler)
+    
+    def checkpoints(self, path):
+        torch.save({
+            "model": self.unet.state_dict(),
+            "optimizer": self.optimizer.state_dict(),
+            "scheduler": self.lr_scheduler.state_dict(),
+            "epoch": self.current_epoch
+        }, path)
+    
+    def load_checkpoint(self, path):
+        checkpoint = torch.load(path)
+        if "model" not in checkpoint:
+            raise ValueError("Checkpoint does not contain a model state dict")
+        else:
+            self.unet.load_state_dict(checkpoint["model"])
+        if "optimizer" in checkpoint:
+            try:
+                self.optimizer.load_state_dict(checkpoint["optimizer"])
+            except:
+                print("Could not load optimizer state dict")
+        if "scheduler" in checkpoint:
+            try:
+                self.lr_scheduler.load_state_dict(checkpoint["scheduler"])
+            except:
+                print("Could not load scheduler state dict")
+        if "epoch" in checkpoint:
+            self.current_epoch = checkpoint["epoch"]
+        else:
+            self.current_epoch = 0
+            print("Epoch information not found in checkpoint setting epoch to 0")
+            self.lr_scheduler.last_epoch = 0
+    
+    def create_random_masks(self, n_masks):
+        masks = torch.ones(n_masks, 1, self.image_size, self.image_size).to(self.device)
+        height = self.masking_range[0][1] - self.masking_range[0][0]
+        width = self.masking_range[1][1] - self.masking_range[1][0]
+        mask_heights = np.random.randint(int(height * self.minimum_mask_portion), int(height * self.maximum_mask_portion), n_masks)
+        mask_widths = np.random.randint(int(width * self.minimum_mask_portion), int(width * self.maximum_mask_portion), n_masks)
+        
+        top_positions = np.random.randint(0, height - mask_heights + 1, n_masks)
+        left_positions = np.random.randint(0, width - mask_widths + 1, n_masks)
+        
+        top_positions += self.masking_range[0][0]
+        left_positions += self.masking_range[1][0]
+        for i in range(n_masks):
+            if np.random.rand() < self.full_image_probability:
+                masks[i] = 0
+                continue
+            top = top_positions[i]
+            left = left_positions[i]
+            mask_height = mask_heights[i]
+            mask_width = mask_widths[i]
+            
+            bottom = top + mask_height
+            right = left + mask_width
+            
+            masks[i, :, top:bottom, left:right] = 0
+            
+        return masks
+    
+    def get_sample_batch(self, n_samples):
+        rnd_idx = np.random.choice(len(self.train_dataset), n_samples).astype(int)
+        images = torch.stack([self.train_dataset[i][0] for i in rnd_idx]).to(self.device)
+        masks = self.create_random_masks(n_samples)
+        masked_images = images * masks + (1-masks)
+        
+        return masked_images.cpu().detach().numpy(), masks.cpu().detach().numpy(), images.cpu().detach().numpy()
+    
+    def get_sample_noising(self, n_samples, n_timesteps=5):
+        rnd_idx = np.random.randint(0, len(self.train_dataset), n_samples)
+        images = torch.stack([self.train_dataset[i][0] for i in rnd_idx])
+        noise = torch.randn(images.shape, device=images.device)
+        timesteps = torch.linspace(0, self.n_train_noise_timesteps-1, n_timesteps, device=images.device).long()
+        
+        noisy_images = torch.zeros((n_samples, n_timesteps, images.shape[1], images.shape[2], images.shape[3]), device=images.device)
+        for i in range(n_timesteps):
+            noisy_images[:, i] = self.noise_scheduler.add_noise(images, noise, timesteps[i])
+        
+        return noisy_images
+        
+    def reset_epoch(self):
+        self.current_epoch = 0
+        self.lr_scheduler.last_epoch = 0
+        
+    def train(self, n_epoch=None, display_interval=1, checkpoint_interval=10, checkpoint_path_prefix="checkpoints"):
+        if n_epoch is None:
+            n_epoch = self.num_epochs
+        
+        accelerator = Accelerator(
+            mixed_precision=self.mixed_precision,
+            gradient_accumulation_steps=self.gradient_accumulation_steps,
+        )
+        
+        
+        model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+            self.unet, self.optimizer, self.dataloader, self.lr_scheduler
+        )
+        
+        for epoch in range(self.current_epoch, n_epoch):
+            model.train()
+            progress = tqdm(train_dataloader)
+            for images in progress: 
+                images = images[0]
+                images = images.to(model.device)
+                bs = images.size(0)
+                # create random masks
+                masks = self.create_random_masks(bs)
+                masked_images = images * masks + (1-masks)
+                
+                noise = torch.randn(images.shape, device=images.device)
+                timesteps = torch.randint(
+                    0, self.noise_scheduler.config.num_train_timesteps, (bs,), device=images.device,
+                    dtype=torch.int64
+                )
+                
+                noisy_images = self.noise_scheduler.add_noise(images, noise, timesteps)
+                
+                with accelerator.accumulate(model):
+                    noise_pred = model(torch.cat([masked_images,noisy_images],1), timesteps, return_dict=False)[0]
+                    
+                    loss = F.mse_loss(noise_pred, noise)
+                    accelerator.backward(loss)
+                    
+                    if accelerator.sync_gradients:
+                        accelerator.clip_grad_norm_(model.parameters(), 1.0)
+                        
+                    optimizer.step()
+                    lr_scheduler.step()
+                    optimizer.zero_grad()
+                
+                progress.set_postfix_str(f"Epoch: {epoch}/{self.num_epochs}, Loss: {loss.item():.7f}")
+            
+            if epoch % display_interval == 0:
+                # unmasks one batch of validation images
+                images = next(iter(self.val_loader))[0].to(model.device)[0:4]
+                masks = self.create_random_masks(images.size(0))
+                masked_images = images * masks + (1-masks)
+                filled_images = self.generate(images=masked_images, num_inference_steps=100)
+                visualize_imagesets(
+                    (images.cpu().detach().numpy() + 1)/2,
+                    masks.cpu().detach().numpy(),
+                    (masked_images.cpu().detach().numpy() + 1)/2,
+                    (filled_images + 1)/2,
+                    titles=["Original", "Mask Maps", "Masked", "Filled"]
+                )
+                
+            if epoch % checkpoint_interval == 0:
+                self.checkpoints(f"{checkpoint_path_prefix}_{epoch}.pt")
+                
+            self.current_epoch += 1
+    
+    def generate(self, images=None, num_inference_steps=100, n_samples=10, noise_seed=None, return_intermediate=False, guidance_function=None, guidance_scale=1.0):
+        if images is None and n_samples is None:
+            raise ValueError("Either images or n_samples must be provided")
+        elif images is None:
+            images = torch.ones(n_samples, 3, self.image_size, self.image_size)
+        
+        initia_noise = None
+        if noise_seed is not None:
+            torch.manual_seed(noise_seed)
+            initia_noise = torch.randn_like(images, device=images.device)
+        
+        self.unet.eval()
+        return self.denoiser(images, num_inference_steps=num_inference_steps, initial_noise=initia_noise, return_intermediate=return_intermediate, guidance_function=guidance_function, guidance_scale=guidance_scale)

bikefusion/pipline.py

@@ -0,0 +1,58 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+from tqdm.auto import tqdm
+
+class InpaintingDenoisingPipeline:
+    def __init__(self, model, scheduler):
+        super().__init__()
+        self.model = model
+        self.scheduler = scheduler
+
+
+    @torch.no_grad()
+    def __call__(self, masked_images, num_inference_steps=50, guidance_function=None, guidance_scale=1.0, initial_noise=None, return_intermediate=False):
+        # Ensure model and scheduler are in evaluation mode
+        self.model.eval()
+
+        # Initialize the noisy image as the input image
+        device = self.model.device
+        masked_images = masked_images.to(device)
+        Bs = masked_images.size(0)
+        
+        if return_intermediate:
+            out = torch.zeros((Bs, num_inference_steps+1, 3, masked_images.shape[2], masked_images.shape[3]), device=device)
+        
+        if initial_noise is not None:
+            noise = torch.tensor(initial_noise).float().to(device)
+        else:
+            noise = torch.randn_like(masked_images, dtype=self.model.dtype).to(self.model.device)
+        
+        self.scheduler.set_timesteps(num_inference_steps)
+        
+        if return_intermediate:
+            out[:, 0] = noise
+            c = 1
+            
+        for t in tqdm(self.scheduler.timesteps):
+            # Get the noise level for this timestep
+            model_output = self.model(torch.cat([masked_images,noise],1), t).sample
+            
+            noise = self.scheduler.step(model_output, t, noise).prev_sample
+            
+            if guidance_function is not None:
+                with torch.enable_grad():
+                    noise.requires_grad = True
+                    guidance_objective = guidance_function(noise, t)
+                    guidance_objective.backward()
+                    grads = noise.grad
+                noise = noise - guidance_scale * grads
+            
+            if return_intermediate:
+                out[:, c] = noise
+                c += 1
+        
+        if return_intermediate:
+            return out.clamp(-1,1).detach().cpu().numpy()
+        else:
+            return noise.clamp(-1,1).detach().cpu().numpy()

bikefusion/visualizers.py

@@ -0,0 +1,88 @@
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+
+def visualize_imagesets(*args,titles=None):
+    """
+    A function to visualize image sets in rows. 
+    Each row corresponds to a different image set with images stacked corresponding to oneanother.
+    
+    Parameters:
+    - args: np.ndarray,torch.tensor the image sets to visualize (must be in B x C x H x W format)
+    """
+    
+    if titles is not None:
+        assert len(titles) == len(args), "Number of titles must match number of image sets"
+    elif titles is None:
+        pass
+    elif not isinstance(titles[0], str) and len(titles[0]) > 1:
+        for i in range(len(args)):
+            assert len(titles[i]) == len(args[i]), "Number of titles must match number of images in each set or be 1"
+    
+    n_sets = len(args)
+    n_images = len(args[0])
+    fig, axs = plt.subplots(n_sets, n_images, figsize=(n_images * 3, n_sets * 3))
+    for i in range(n_sets):
+        for j in range(n_images):
+            if isinstance(args[i][j], torch.Tensor):
+                image = args[i][j].detach().cpu().numpy()
+            elif isinstance(args[i][j], np.ndarray):
+                image = args[i][j]
+            else:
+                raise ValueError("Image must be a numpy array or torch tensor")
+            if image.shape[0] == 1:
+                if n_sets == 1:
+                    axs[j].imshow(image[0], cmap="gray")
+                else:
+                    axs[i, j].imshow(image[0], cmap="gray")
+            else:
+                if n_sets == 1:
+                    axs[j].imshow(np.transpose(image, (1, 2, 0)))
+                else:
+                    axs[i, j].imshow(np.transpose(image, (1, 2, 0)))
+            if n_sets == 1:
+                axs[j].axis("off")
+            else:
+                axs[i, j].axis("off")
+            if titles is not None and not isinstance(titles[i], str):
+                if len(titles[i]) == 1:
+                    axs[j].set_title(titles[i][0])
+                else:
+                    axs[i, j].set_title(titles[i][j])
+            elif titles is not None:
+                if n_sets == 1:
+                    axs[j].set_title(titles[i])
+                else:
+                    axs[i, j].set_title(titles[i])
+                
+                
+    plt.tight_layout()
+    plt.show()
+    
+def visualize_image_evolution(images, titles=None):
+    """
+    A function to visualize the evolution of images over time.
+    
+    Parameters:
+    - images: np.ndarray, shape (n_images, n_timesteps, channels, height, width), the images to visualize
+    - titles: list, the titles for each iteration
+    """
+    if isinstance(images, torch.Tensor):
+        images = images.detach().cpu().numpy()
+        
+    if titles is not None:
+        assert len(titles) == images.shape[1], "Number of titles must match number of timesteps"
+    n_images = images.shape[0]
+    n_timesteps = images.shape[1]
+    fig, axs = plt.subplots(n_images, n_timesteps, figsize=(n_timesteps * 3, n_images * 3))
+    for i in range(n_images):
+        for j in range(n_timesteps):
+            if images.shape[2] == 1:
+                axs[i, j].imshow(images[i, j, 0], cmap="gray")
+            else:
+                axs[i, j].imshow(np.transpose(images[i, j], (1, 2, 0)))
+            axs[i, j].axis("off")
+            if titles is not None:
+                axs[i, j].set_title(titles[j])
+    plt.tight_layout()
+    plt.show()

chekpoint/bikefusion.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8e68138c27aea63527caa1344ed9e1aa9ad21580f0ac8194b4495917b20409ee
+size 678272190

data/desc_test.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2afd555e135b03e8176e726e692d26c131838bf6ba397c13b1238582401e21c7
+size 682467

data/desc_train.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8314c63c08692c3f016b7bdd87d97620e4e978aa72cbd637f79ef4f6438babf2
+size 6844466

data/images_train_0.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d003aba1dd5d8a35b7b467ef926d09a7b1d4ed4e309876ed49f5f81c181125c9
+size 61440128

data/images_train_1.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34560a1a54310f8dd8ca831f2510df098323ef2c7509e2aaeed6f80d319f46a1
+size 61440128

data/images_train_2.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8035d533235b5af64352964181a2a2d8ffb051de7def627d4cf6ceba641e04b
+size 61440128

data/images_train_3.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba7a6d2c142189b1278aa7fd528d90af054531884acd923747a66f0f4a9668ca
+size 61440128

data/images_train_4.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e23641baa6c970ebb16c2c1a0b34d4d63201fba7d03028673de4c8ee466ea0e
+size 61440128

data/mask_test.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:341234393c5924a807c991ef17f788c1a559ecb37d0e43e515bb372f7d9d526d
+size 32128

data/mask_train.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b01ffbdc95fea079ddc07151ea3a7eb282aea5b38467dbe738c4e420973a9b8f
+size 320128

data/masked_test.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52cb7f4d6b14828ad8bafc095ebf4ec4fa33aeda7d8a8a7e112365b279b49e67
+size 30720128

data/masked_train_0.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32183c7f1c0600392cc4f16429162c1982e2495241036eb5576af3db2d49e50d
+size 61440128

data/masked_train_1.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5f5f32b54c71830452fef97e33fad31d8be6133cd99988f59639de94a3b5a2c
+size 61440128

data/masked_train_2.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:971dd8d938d8d29cfbdd7b74811b6fa6870646795601768c6bf9e79bf85be61c
+size 61440128

data/masked_train_3.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6bf79d2d514de1840f7b583d91164abb75f49bc0470547da9609027be84cdcc5
+size 61440128

data/masked_train_4.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e7eaaadb2a3673d9f95639a891bf53d0c3dcb8fbfeed9e011720de6a9dbd51c5
+size 61440128

data/param_test.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7eea53c5db5f0062570ae63deacc0df05ccada087781940c648cd54d3f5390df
+size 768128

data/param_train.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d267b272bf69c08cc302190bbf47eb7a0e31edc70180cf7010d02888f2b1e8fb
+size 7680128