Views
No views yet
1from transformers import AutoModel
2import torch
3import torch.nn as nn
4import torchvision
5from torchvision import transforms as v2
6import numpy as np
7
8# Noise Injector transformation
9class SaturationNoiseInjector(nn.Module):
10 def __init__(self, low=200, high=255):
11 super().__init__()
12 self.low = low
13 self.high = high
14
15 def forward(self, x: torch.Tensor) -> torch.Tensor:
16 channel = x[0].clone()
17 noise = torch.empty_like(channel).uniform_(self.low, self.high)
18 mask = (channel == 255).float()
19 noise_masked = noise * mask
20 channel[channel == 255] = 0
21 channel = channel + noise_masked
22 x[0] = channel
23 return x
24
25
26# Self Normalize transformation
27class PerImageNormalize(nn.Module):
28 def __init__(self, eps=1e-7):
29 super().__init__()
30 self.eps = eps
31 self.instance_norm = nn.InstanceNorm2d(
32 num_features=1,
33 affine=False,
34 track_running_stats=False,
35 eps=self.eps,
36 )
37
38 def forward(self, x: torch.Tensor) -> torch.Tensor:
39 if x.dim() == 3:
40 x = x.unsqueeze(0)
41 x = self.instance_norm(x)
42 if x.shape[0] == 1:
43 x = x.squeeze(0)
44 return x
45
46
47# Load model
48device = "cuda" if torch.cuda.is_available() else "cpu"
49model = AutoModel.from_pretrained("CaicedoLab/MorphEm", trust_remote_code=True)
50model.to(device).eval()
51
52# Define transforms
53transform = v2.Compose([
54 SaturationNoiseInjector(),
55 PerImageNormalize(),
56 v2.Resize(size=(224, 224), antialias=True),
57])
58
59# Generate random batch (N, C, H, W)
60batch_size = 2
61num_channels = 3
62images = torch.randint(0, 256, (batch_size, num_channels, 512, 512), dtype=torch.float32)
63
64print(f"Input shape: {images.shape} (N={batch_size}, C={num_channels}, H=512, W=512)")
65print()
66
67# Bag of Channels (BoC) - process each channel independently
68with torch.no_grad():
69 batch_feat = []
70 images = images.to(device)
71
72 for c in range(images.shape[1]):
73 # Extract single channel: (N, C, H, W) -> (N, H, W)
74 single_channel = images[:, c, :, :]
75
76 # Apply transforms, add dimension 1 ((N, 1, H, W))
77 single_channel = transform(single_channel).unsqueeze(1)
78
79 # Extract features
80 output = model.forward_features(single_channel)
81 feat_temp = output["x_norm_clstoken"].cpu().detach().numpy()
82 batch_feat.append(feat_temp)
83
84# Concatenate features from all channels
85features = np.concatenate(batch_feat, axis=1)
86
87print(f"Output shape: {features.shape}")
88print(f" - Batch size (N): {features.shape[0]}")
89print(f" - Feature dimension (C * feature_dim): {features.shape[1]}")1# Noise Injector transformation
2class SaturationNoiseInjector(nn.Module):
3 def __init__(self, low=200, high=255):
4 super().__init__()
5 self.low = low
6 self.high = high
7
8 def forward(self, x: torch.Tensor) -> torch.Tensor:
9 channel = x[0].clone()
10 noise = torch.empty_like(channel).uniform_(self.low, self.high)
11 mask = (channel == 255).float()
12 noise_masked = noise * mask
13 channel[channel == 255] = 0
14 channel = channel + noise_masked
15 x[0] = channel
16 return x
17
18
19# Self Normalize transformation
20class PerImageNormalize(nn.Module):
21 def __init__(self, eps=1e-7):
22 super().__init__()
23 self.eps = eps
24 self.instance_norm = nn.InstanceNorm2d(
25 num_features=1,
26 affine=False,
27 track_running_stats=False,
28 eps=self.eps,
29 )
30
31 def forward(self, x: torch.Tensor) -> torch.Tensor:
32 if x.dim() == 3:
33 x = x.unsqueeze(0)
34 x = self.instance_norm(x)
35 if x.shape[0] == 1:
36 x = x.squeeze(0)
37 return x1@inproceedings{
2agrawal2026chammi,
3title={{CHAMMI}-75: Pre-training multi-channel models with heterogeneous microscopy images},
4author={Vidit Agrawal and John Peters and Tyler N. Thompson and Mohammad Vali Sanian and Chau Pham and Nikita Moshkov and Arshad Kazi and Aditya Pillai and Jack Freeman and Byunguk Kang and Samouil L. Farhi and Ernest Fraenkel and Ron M. Stewart and Lassi Paavolainen and Bryan A. Plummer and Juan C. Caicedo},
5booktitle={The Fourteenth International Conference on Learning Representations},
6year={2026},
7url={https://openreview.net/forum?id=SLjqdj3LPk}
8}