1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
|
"""
Utility functions for MNIST diffusion.
"""
import os
import io
import numpy as np
import torch
from PIL import Image
import plotly.graph_objects as go
from plotly.subplots import make_subplots
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: go.Figure) -> np.ndarray:
"""
Convert a plotly figure to a numpy image array.
Args:
figure: Plotly figure
Returns:
Numpy array of the figure image
"""
buf = io.BytesIO()
figure.write_image(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: go.Figure, filename: str):
"""
Save a plotly figure to the plots directory.
Args:
figure: Plotly figure to save
filename: Filename (will be saved in PLOTS_DIR)
"""
ensure_dirs()
path = os.path.join(PLOTS_DIR, filename)
figure.write_image(path)
|
