trying to improve whole project

1 files changed, 110 insertions, 0 deletions

diff --git a/src/utils.py b/src/utils.py
new file mode 100644
index 0000000..f727a68
--- /dev/null
+++ b/src/utils.py
@@ -0,0 +1,110 @@
+"""
+Utility functions for MNIST diffusion.
+"""
+
+import os
+import io
+import numpy as np
+import torch
+import scipy.linalg
+from PIL import Image
+import matplotlib.pyplot as plt
+
+from .config import CHECKPOINT_DIR, PLOTS_DIR
+
+
+def ensure_dirs():
+    """Create necessary directories if they don't exist."""
+    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+    os.makedirs(PLOTS_DIR, exist_ok=True)
+
+
+def tensor_to_image(tensor: torch.Tensor) -> np.ndarray:
+    """
+    Convert a tensor to a numpy image array.
+
+    Args:
+        tensor: Image tensor [C, H, W]
+
+    Returns:
+        Numpy array [H, W, C] normalized to [0, 1]
+    """
+    img = tensor.clone().detach().cpu().numpy().transpose(1, 2, 0)
+    img -= img.min()
+    img /= img.max() + 1e-8
+    return img
+
+
+def tensor_to_images(tensor: torch.Tensor) -> np.ndarray:
+    """
+    Convert a batch of tensors to numpy image arrays.
+
+    Args:
+        tensor: Batch of image tensors [B, C, H, W]
+
+    Returns:
+        Numpy array [B, H, W, C] normalized to [0, 1]
+    """
+    img = tensor.clone().detach().cpu().numpy().transpose(0, 2, 3, 1)
+    img -= np.min(img, axis=(1, 2, 3), keepdims=True)
+    img /= np.max(img, axis=(1, 2, 3), keepdims=True) + 1e-8
+    return img
+
+
+def figure_to_image(figure: plt.Figure) -> np.ndarray:
+    """
+    Convert a matplotlib figure to a numpy image array.
+
+    Args:
+        figure: Matplotlib figure
+
+    Returns:
+        Numpy array of the figure image
+    """
+    buf = io.BytesIO()
+    figure.savefig(buf, format='png')
+    buf.seek(0)
+    image = np.array(Image.open(buf))
+    return image
+
+
+def save_checkpoint(model: torch.nn.Module, filename: str):
+    """
+    Save model checkpoint.
+
+    Args:
+        model: Model to save
+        filename: Filename (will be saved in CHECKPOINT_DIR)
+    """
+    ensure_dirs()
+    path = os.path.join(CHECKPOINT_DIR, filename)
+    torch.save(model.state_dict(), path)
+
+
+def load_checkpoint(model: torch.nn.Module, filename: str) -> torch.nn.Module:
+    """
+    Load model checkpoint.
+
+    Args:
+        model: Model architecture to load weights into
+        filename: Filename (loaded from CHECKPOINT_DIR)
+
+    Returns:
+        Model with loaded weights
+    """
+    path = os.path.join(CHECKPOINT_DIR, filename)
+    model.load_state_dict(torch.load(path, weights_only=True))
+    return model
+
+
+def save_plot(figure: plt.Figure, filename: str):
+    """
+    Save a matplotlib figure to the plots directory.
+
+    Args:
+        figure: Matplotlib figure to save
+        filename: Filename (will be saved in PLOTS_DIR)
+    """
+    ensure_dirs()
+    path = os.path.join(PLOTS_DIR, filename)
+    figure.savefig(path)