deleting unecessary files/scripts

config/confg.yaml

@@ -1,64 +0,0 @@
-data_confg:
-  train_batch_size: 6
-  val_batch_size: 6
-  test_batch_size: 6
-  flow_method: raft
-  data_domain: animation
-  datamodule_confg:
-    mean: [0.5, 0.5, 0.5]
-    sd: [0.5, 0.5, 0.5]
-    size: [256, 448]
-    amount_augmentations: 1
-    horizontal_flip: 0.5
-    time_flip: True
-    rotation: 0
-    brightness: 0.2
-    contrast: 0.2
-    saturation: 0.2
-    hue: 0.1
-
-trainer_confg:
-  accumulate_grad_batches: 5
-  gradient_clip_val: 1.0
-  max_epochs: 500
-  num_nodes: 1
-  devices: 2
-  accelerator: gpu
-  strategy: ddp_find_unused_parameters_true
-
-optim_confg:
-  optimizer_confg: # AdamW
-    lr: 1.0e-4
-    betas: [0.9, 0.999]
-    eps: 1.0e-8
-  scheduler_confg: # ReduceLROnPlateau
-    mode: min
-    factor: 0.5
-    patience: 3
-    verbose: True
-
-pretrained_model_path: null # Fine-tune model path
-
-model_confg:
-  kappa: 2.0
-  timesteps: 20
-  p: 0.3
-  etas_end: 0.99
-  min_noise_level: 0.04
-  flow_model: raft
-  flow_kwargs:
-    pretrained_path: null #_pretrain_models/anime_interp_full.ckpt
-  warping_kwargs:
-    in_channels: 3
-    channels: [128, 256, 384, 512]
-  synthesis_kwargs:
-    in_channels: 3
-    channels: [128, 256, 384, 512]
-    temb_channels: 512
-    heads: 1
-    window_size: 8
-    window_attn: True
-    grid_attn: True
-    expansion_rate: 1.5
-    num_conv_blocks: 1
-    dropout: 0.0

model/train_pipline.py

@@ -1,177 +0,0 @@
-import os
-import copy
-import matplotlib.pyplot as plt
-from typing import Any
-
-import torch
-from torch.optim.lr_scheduler import ReduceLROnPlateau
-from torch.optim import AdamW, Optimizer
-from torch.utils.data import DataLoader
-from lightning import LightningModule
-
-from torchmetrics import MetricCollection
-from torchmetrics.image import PeakSignalNoiseRatio as PSNR
-from torchmetrics.image import StructuralSimilarityIndexMeasure as SSIM
-from torchmetrics.image import LearnedPerceptualImagePatchSimilarity as LPIPS
-
-from model.model import MultiInputResShift
-
-from utils.utils import denorm, make_grid_images#, save_triplet 
-from utils.ema import EMA
-from utils.inter_frame_idx import get_inter_frame_temp_index
-from utils.raft import raft_flow
-
-
-class TrainPipline(LightningModule):
-    def __init__(self, 
-                 confg: dict, 
-                 test_dataloader: DataLoader):
-        super(TrainPipline, self).__init__()
-
-        self.test_dataloader = test_dataloader
-
-        self.confg = confg
-
-        self.mean, self.sd = confg["data_confg"]["mean"], confg["data_confg"]["sd"]
-
-        self.model = MultiInputResShift(**confg["model_confg"])
-        self.model.flow_model.requires_grad_(False).eval()
-
-        self.ema = EMA(beta=0.995)
-        self.ema_model = copy.deepcopy(self.model).eval().requires_grad_(False)
-
-        self.charbonnier_loss = lambda x, y: torch.mean(torch.sqrt((x - y)**2 + 1e-6))
-        self.lpips_loss = LPIPS(net_type='vgg')
-
-        self.train_metrics = MetricCollection({
-            "train_lpips": LPIPS(net_type='alex'),
-            "train_psnr": PSNR(),
-            "train_ssim": SSIM()
-        })
-        self.val_metrics = MetricCollection({
-            "val_lpips": LPIPS(net_type='alex'),
-            "val_psnr": PSNR(),
-            "val_ssim": SSIM()
-        })
-
-    def loss_fn(self, 
-                x: torch.Tensor, 
-                predicted_x: torch.Tensor) -> torch.Tensor:
-        percep_loss = 0.2 * self.lpips_loss(x, predicted_x.clamp(-1, 1))
-        pix2pix_loss = self.charbonnier_loss(x, predicted_x)
-        return percep_loss + pix2pix_loss
-
-    def sample_t(self, 
-                 shape: tuple[int, ...], 
-                 max_t: int, 
-                 device: torch.device) -> torch.Tensor:
-        p = torch.linspace(1, max_t, steps=max_t, device=device) ** 2 
-        p = p / p.sum()  
-        t = torch.multinomial(p, num_samples=shape[0], replacement=True)
-        return t
-
-    def forward(self, 
-                I0: torch.Tensor, 
-                It: torch.Tensor, 
-                I1: torch.Tensor) -> torch.Tensor:
-        flow0tot = raft_flow(I0, It, 'animation')
-        flow1tot = raft_flow(I1, It, 'animation')
-        mid_idx = get_inter_frame_temp_index(I0, It, I1, flow0tot, flow1tot).to(It.dtype)
-
-        tau = torch.stack([mid_idx, 1 - mid_idx], dim=1)
-        
-        if self.current_epoch > 5:
-            t = torch.randint(low=1, high=self.model.timesteps, size=(It.shape[0],), device=It.device, dtype=torch.long)
-        else:
-            t = self.sample_t(shape=(It.shape[0],), max_t=self.model.timesteps, device=It.device)
-
-        predicted_It = self.model(I0, It, I1, tau=tau, t=t)
-        return predicted_It
-
-    def get_step_plt_images(self, 
-                            It: torch.Tensor, 
-                            predicted_It: torch.Tensor) -> plt.Figure:
-        fig, ax = plt.subplots(1, 2, figsize=(20, 10))
-        ax[0].imshow(denorm(predicted_It.clamp(-1, 1), self.mean, self.sd)[0].permute(1, 2, 0).cpu().numpy())
-        ax[0].axis("off")
-        ax[0].set_title("Predicted")
-        ax[1].imshow(denorm(It, self.mean, self.sd)[0].permute(1, 2, 0).cpu().numpy())
-        ax[1].axis("off")
-        ax[1].set_title("Ground Truth")
-        plt.tight_layout() 
-        #img_path = "step_image.png"
-        #fig.savefig(img_path, dpi=300, bbox_inches='tight') 
-        plt.close(fig)
-        return fig
-
-    def training_step(self, batch: tuple[torch.Tensor, ...], _) -> torch.Tensor:
-        I0, It, I1 = batch
-        predicted_It = self(I0, It, I1)
-        loss = self.loss_fn(It, predicted_It)
-
-        self.log("lr", self.trainer.optimizers[0].param_groups[0]["lr"], prog_bar=True, on_step=True, on_epoch=False, sync_dist=True)
-        self.log("train_loss", loss, prog_bar=True, on_step=True, on_epoch=False, sync_dist=True)
-
-        self.ema.step_ema(self.ema_model, self.model)
-        with torch.inference_mode():
-            fig = self.get_step_plt_images(It, predicted_It)
-            self.logger.experiment.add_figure("Train Predictions", fig, self.global_step)
-            mets = self.train_metrics(It, predicted_It.clamp(-1, 1))
-            self.log_dict(mets, prog_bar=True, on_step=True,on_epoch=False)
-        return loss
-    
-    @torch.no_grad()
-    def validation_step(self,  batch: tuple[torch.Tensor, ...], _) -> None:
-        I0, It, I1 = batch
-        predicted_It = self(I0, It, I1)
-        loss = self.loss_fn(It, predicted_It)
-        
-        self.log("val_loss", loss, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True)
-
-        mets = self.val_metrics(It, predicted_It.clamp(-1, 1))
-        self.log_dict(mets, prog_bar=True, on_step=False, on_epoch=True)
-
-    @torch.inference_mode()
-    def on_train_epoch_end(self) -> None:
-        torch.save(self.ema_model.state_dict(), 
-                   os.path.join("_checkpoint", f"resshift_diff_{self.current_epoch}.pth"))
-
-        batch = next(iter(self.test_dataloader))
-        I0, It, I1 = batch
-        I0, It, I1 = I0.to(self.device), It.to(self.device), I1.to(self.device)
-
-        flow0tot = raft_flow(I0, It, 'animation')
-        flow1tot = raft_flow(I1, It, 'animation')
-        mid_idx = get_inter_frame_temp_index(I0, It, I1, flow0tot, flow1tot).to(It.dtype)
-        tau = torch.stack([mid_idx, 1 - mid_idx], dim=1)
-
-        predicted_It = self.ema_model.reverse_process([I0, I1], tau)
-
-        I0 = denorm(I0, self.mean, self.sd)
-        I1 = denorm(I1, self.mean, self.sd)
-        It = denorm(It, self.mean, self.sd)
-        predicted_It = denorm(predicted_It.clamp(-1, 1), self.mean, self.sd)
-
-        #save_triplet([I0, It, predicted_It, I1], f"./_output/target_{self.current_epoch}.png", nrow=1)
-        grid = make_grid_images([I0, It, predicted_It, I1], nrow=1)
-        self.logger.experiment.add_image("Predicted Images", grid, self.global_step)
-
-    def configure_optimizers(self) -> tuple[list[Optimizer], list[dict[str, Any]]]:
-        optimizer = [AdamW(
-                        self.model.parameters(),
-                        **self.confg["optim_confg"]['optimizer_confg']
-                    )]
-         
-        scheduler = [{
-            'scheduler': ReduceLROnPlateau(
-                optimizer[0],
-                **self.confg["optim_confg"]['scheduler_confg']
-            ),
-            'monitor': 'val_loss',
-            'interval': 'epoch',
-            'frequency': 1,
-            'strict': True,
-        }] 
-                    
-        return optimizer, scheduler
-

utils/ema.py

@@ -1,32 +0,0 @@
-import torch
-import torch.nn as nn
-
-class EMA:
-    def __init__(self, beta: float):
-        super().__init__()
-        self.beta = beta
-        self.step = 0
-
-    def update_model_average(self, ema_model: nn.Module, current_model: nn.Module) -> None:
-        for current_params, ema_model in zip(current_model.parameters(), ema_model.parameters()):
-            old_weight, up_weight = ema_model.data, current_params.data
-            ema_model.data = self.update_average(old_weight, up_weight)
-
-    def update_average(self, old: torch.Tensor | None, new: torch.Tensor) -> torch.Tensor:
-        if old is None:
-            return new
-        return old * self.beta + (1 - self.beta) * new
-
-    def step_ema(self, ema_model: nn.Module, model: nn.Module, step_start_ema: int = 2000) -> None:
-        if self.step < step_start_ema:
-            self.reset_parameters(ema_model, model)
-            self.step += 1
-            return
-        self.update_model_average(ema_model, model)
-        self.step += 1
-
-    def copy_to(self, ema_model: nn.Module, model: nn.Module) -> None:
-        model.load_state_dict(ema_model.state_dict())
-
-    def reset_parameters(self, ema_model: nn.Module, model: nn.Module) -> None:
-        ema_model.load_state_dict(model.state_dict())

utils/inter_frame_idx.py

@@ -1,123 +0,0 @@
-from utils.utils import morph_open
-
-import torch
-from kornia.color import rgb_to_grayscale
-
-import cv2
-import numpy as np
-
-class FlowEstimation:
-    def __init__(self, flow_estimator: str = "farneback"):
-        assert flow_estimator in ["farneback", "dualtvl1"], "Flow estimator must be one of [farneback, dualtvl1]"
-
-        if flow_estimator == "farneback":
-            self.flow_estimator = self.OptFlow_Farneback
-        elif flow_estimator == "dualtvl1":
-            self.flow_estimator = self.OptFlow_DualTVL1
-        else:
-            raise NotImplementedError
-
-    def OptFlow_Farneback(self, I0: torch.Tensor, I1: torch.Tensor) -> torch.Tensor:
-        device = I0.device
-        
-        I0 = I0.cpu().clamp(0, 1) * 255
-        I1 = I1.cpu().clamp(0, 1) * 255
-
-        batch_size = I0.shape[0]
-        for i in range(batch_size):
-            I0_np = I0[i].permute(1, 2, 0).numpy().astype(np.uint8)
-            I1_np = I1[i].permute(1, 2, 0).numpy().astype(np.uint8)
-
-            I0_gray = cv2.cvtColor(I0_np, cv2.COLOR_BGR2GRAY)
-            I1_gray = cv2.cvtColor(I1_np, cv2.COLOR_BGR2GRAY)
-
-            flow = cv2.calcOpticalFlowFarneback(I0_gray, I1_gray, None, 0.5, 3, 15, 3, 5, 1.2, 0)
-            flow = torch.from_numpy(flow).permute(2, 0, 1).unsqueeze(0).float()
-            if i == 0:
-                flows = flow
-            else:
-                flows = torch.cat((flows, flow), dim = 0)
-
-        return flows.to(device)
-
-    def OptFlow_DualTVL1(
-        self, 
-        I0: torch.Tensor, 
-        I1: torch.Tensor,
-        tau: float = 0.25,
-        lambda_: float = 0.15,
-        theta: float = 0.3,
-        scales_number: int = 5,
-        warps: int = 5,
-        epsilon: float = 0.01,
-        inner_iterations: int = 30,
-        outer_iterations: int = 10,
-        scale_step: float = 0.8,
-        gamma: float = 0.0
-    ) -> torch.Tensor:
-        optical_flow = cv2.optflow.createOptFlow_DualTVL1()
-        optical_flow.setTau(tau)
-        optical_flow.setLambda(lambda_)
-        optical_flow.setTheta(theta)
-        optical_flow.setScalesNumber(scales_number)
-        optical_flow.setWarpingsNumber(warps)
-        optical_flow.setEpsilon(epsilon)
-        optical_flow.setInnerIterations(inner_iterations)
-        optical_flow.setOuterIterations(outer_iterations)
-        optical_flow.setScaleStep(scale_step)
-        optical_flow.setGamma(gamma)
-
-        device = I0.device
-        
-        I0 = I0.cpu().clamp(0, 1) * 255
-        I1 = I1.cpu().clamp(0, 1) * 255
-
-        batch_size = I0.shape[0]
-        for i in range(batch_size):
-            I0_np = I0[i].permute(1, 2, 0).numpy().astype(np.uint8)
-            I1_np = I1[i].permute(1, 2, 0).numpy().astype(np.uint8)
-
-            I0_gray = cv2.cvtColor(I0_np, cv2.COLOR_BGR2GRAY)
-            I1_gray = cv2.cvtColor(I1_np, cv2.COLOR_BGR2GRAY)
-
-            flow = optical_flow.calc(I0_gray, I1_gray, None)
-            flow = torch.from_numpy(flow).permute(2, 0, 1).unsqueeze(0).float()
-            if i == 0:
-                flows = flow
-            else:
-                flows = torch.cat((flows, flow), dim = 0)
-
-        return flows.to(device)
-    
-    def __call__(self, I1: torch.Tensor, I0: torch.Tensor) -> torch.Tensor:
-        return self.flow_estimator(I1, I0)
-
-def get_inter_frame_temp_index(
-    I0: torch.Tensor, 
-    It: torch.Tensor, 
-    I1: torch.Tensor, 
-    flow0tot: torch.Tensor, 
-    flow1tot: torch.Tensor, 
-    k: int = 5, 
-    threshold: float = 2e-2
-) -> torch.Tensor:
-
-    I0_gray = rgb_to_grayscale(I0)
-    It_gray = rgb_to_grayscale(It)
-    I1_gray = rgb_to_grayscale(I1)
-
-    mask0tot = morph_open(It_gray - I0_gray, k=k)
-    mask1tot = morph_open(I1_gray - It_gray, k=k)
-
-    mask0tot = (abs(mask0tot) > threshold).to(torch.uint8)
-    mask1tot = (abs(mask1tot) > threshold).to(torch.uint8)
-
-    flow_mag0tot = torch.sqrt(flow0tot[:, 0, :, :]**2 + flow0tot[:, 1, :, :]**2).unsqueeze(1)
-    flow_mag1tot = torch.sqrt(flow1tot[:, 0, :, :]**2 + flow1tot[:, 1, :, :]**2).unsqueeze(1)
-
-    norm0tot = (flow_mag0tot*mask0tot).squeeze(1)
-    norm1tot = (flow_mag1tot*mask1tot).squeeze(1)
-    d0tot = torch.sum(norm0tot, dim = (1, 2)) 
-    d1tot = torch.sum(norm1tot, dim = (1, 2))
-    
-    return d0tot / (d0tot + d1tot + 1e-12)

utils/raft.py

@@ -1,20 +0,0 @@
-import torch
-from torchvision.models.optical_flow import raft_large 
-from modules.flow_models.raft.rfr_new import RAFT
-
-def raft_flow(
-        I0: torch.Tensor, 
-        I1: torch.Tensor, 
-        data_domain: str = "animation", 
-        device: str = 'cuda'
-) -> tuple[torch.Tensor, torch.Tensor]:
-    if I0.dtype != torch.float32 or I1.dtype != torch.float32:
-        I0 = I0.to(torch.float32)
-        I1 = I1.to(torch.float32)
-    if data_domain == "animation":
-        raft = RAFT().requires_grad_(False).eval().to(device)
-    elif data_domain == "photorealism":
-        raft = raft_large().requires_grad_(False).eval().to(device)
-    else:
-        raise ValueError("data_domain must be either 'animation' or 'photorealism'")
-    return raft(I0, I1) if data_domain == "animation" else raft(I0, I1)[-1]

utils/uncertainty.py

@@ -1,49 +0,0 @@
-import torch
-import itertools
-from torchmetrics.image import LearnedPerceptualImagePatchSimilarity as LPIPS
-from utils.utils import denorm
-
-def compute_lpips_variability(samples: torch.Tensor, 
-                              net: str = 'alex', 
-                              device: str = 'cuda'
-                              ) -> float:
-    loss_fn = LPIPS(net_type=net).to(device)
-    loss_fn.eval()
-
-    if samples.min() >= 0.0:
-        samples = samples * 2 - 1  # Convertir [0, 1] → [-1, 1]
-
-    N = samples.size(0)
-    scores = []
-    for i, j in itertools.combinations(range(N), 2):
-        x = samples[i:i+1].to(device)
-        y = samples[j:j+1].to(device)
-        dist = loss_fn(denorm(x.clamp(-1, 1)), denorm(y.clamp(-1, 1)))
-        scores.append(dist.item())
-
-    return sum(scores) / len(scores)
-
-def compute_pixelwise_correlation(samples: torch.Tensor) -> float:
-    N, C, H, W = samples.shape
-    samples_flat = samples.view(N, C, -1)  # (N, C, H*W)
-
-    corrs = []
-    for i, j in itertools.combinations(range(N), 2):
-        x = samples_flat[i]  # (C, HW)
-        y = samples_flat[j]  # (C, HW)
-        mean_x = x.mean(dim=1, keepdim=True)
-        mean_y = y.mean(dim=1, keepdim=True)
-        x_centered = x - mean_x
-        y_centered = y - mean_y
-        numerator = (x_centered * y_centered).sum(dim=1)
-        denominator = (x_centered.norm(dim=1) * y_centered.norm(dim=1)) + 1e-8
-        corr = numerator / denominator  # (C,)
-        corrs.append(corr.mean().item())
-    return sum(corrs) / len(corrs)
-
-def compute_dynamic_range(samples: torch.Tensor) -> float:
-    max_vals, _ = samples.max(dim=0)  # (C, H, W)
-    min_vals, _ = samples.min(dim=0)  # (C, H, W)
-    
-    dynamic_range = max_vals - min_vals  # (C, H, W)
-    return dynamic_range.mean().item()