updated load scripts

__init__.py

@@ -1,23 +1,8 @@
 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)
+load_bikefusion = lambda: load_bikefusion_and_data(current_dir)
 
-    # 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
+__all__ = [name for name in dir() if not name.startswith("_") and name not in {"os"}]

bikefusion/__init__.py

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

bikefusion/data_utils.py

@@ -2,6 +2,9 @@ import torchvision.transforms.functional as F
 import torch
 import numpy as np
 import pickle
+from torch.utils.data import TensorDataset, random_split
+from .diffusion import InpaintingDenoisingDiffusion
+import os
 
 def to_mask_map(maks, image_size=(80,128)):
     bs = maks.shape[0]
@@ -96,4 +99,21 @@ def load_data(split="train", path="data/"):
     masks = np.load(f"{path}mask_{split}.npy")
 
 
-    return masked_images, masks, parametric, description, images
+    return masked_images, masks, parametric, description, images
+
+def load_bikefusion_and_data(current_dir):
+    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