"""Sample a random bounding box for window cropping.

W = size[2]

H = size[3]

while True:

if method == "original":

ratio = np.random.beta(beta, beta)

cut_rat = np.sqrt(ratio)

cut_w = int(W * cut_rat)

cut_h = int(H * cut_rat)

cx = np.random.randint(W)

cy = np.random.randint(H)

bbx1 = int(np.clip(cx - cut_w // 2, 0, W))

bby1 = int(np.clip(cy - cut_h // 2, 0, H))

bbx2 = int(np.clip(cx + cut_w // 2, 0, W))

bby2 = int(np.clip(cy + cut_h // 2, 0, H))

elif method == "fixedShape":

ratio = np.random.beta(beta, beta)

cut_rat = np.sqrt(ratio)

cut_w = int(W * cut_rat)

cut_h = int(H * cut_rat)

bbx1 = np.random.randint(0, W - cut_w)

bby1 = np.random.randint(0, H - cut_h)

bbx2 = bbx1 + cut_w

bby2 = bby1 + cut_h

elif method == "randomShape":

scale = np.random.beta(beta, beta)

while True:

ratio = np.random.uniform(0.3, 1 / 0.3)

w_rat = np.sqrt(scale * ratio)

h_rat = np.sqrt(scale / ratio)

cut_w = int(W * w_rat)

cut_h = int(H * h_rat)

if W - cut_w > 0 and H - cut_h > 0:

break

bbx1 = np.random.randint(0, W - cut_w)

bby1 = np.random.randint(0, H - cut_h)

bbx2 = bbx1 + cut_w

bby2 = bby1 + cut_h

elif method == "fixedCenter":

ratio = np.random.beta(beta, beta)

cut_rat = np.sqrt(ratio)

cut_w = int(W * cut_rat)

cut_h = int(H * cut_rat)

cx = W // 2

cy = H // 2

bbx1 = int(np.clip(cx - cut_w // 2, 0, W))

bby1 = int(np.clip(cy - cut_h // 2, 0, H))

bbx2 = int(np.clip(cx + cut_w // 2, 0, W))

bby2 = int(np.clip(cy + cut_h // 2, 0, H))

elif method == "vertex":

ratio = np.random.beta(beta, beta)

cut_rat = np.sqrt(ratio)

cut_w = int(W * cut_rat)

cut_h = int(H * cut_rat)

bbx1 = np.random.choice([0, W - cut_w])

bby1 = np.random.choice([0, H - cut_h])

bbx2 = bbx1 + cut_w

bby2 = bby1 + cut_h

else:

raise ValueError(f"Unknown bbox sampling method: '{method}'")

area_ratio = float(bbx2 - bbx1) * (bby2 - bby1) / (W * H)

if area_ratio >= bbx_thres:

break

"""Compute per-instance, per-channel standard deviation and mean.

N, C = x.size()[:2]

if x.dim() == 4:

std, mean = torch.std_mean(x, dim=(2, 3), keepdim=True)

elif x.dim() == 3:

std, mean = torch.std_mean(x, dim=2, keepdim=True)

else:

raise NotImplementedError(

f"Unsupported feature shape with {x.dim()} dimensions."

)

"""Window Normalization (WIN) for 2-D feature maps.

# Class-level bbox cache shared across all instances to avoid redundant

# disk I/O. All WIN layers in a single model share the same pre-computed

# bbox pool — this is intentional.

_cached_bboxs = None

_cached_bboxs_len = 0

def __init__(

self,

num_features,

mask_thres=0.7,

eps=1e-5,

alpha=0.1,

mix=True,

grid=False,

input_size=224,

mask_patch_size=32,

affine=False,

cached=True,

bboxs_path="./bboxs.npy",

device="cuda",

):

super().__init__()

if cached and WindowNorm2d._cached_bboxs is None:

WindowNorm2d._cached_bboxs = np.load(bboxs_path)

WindowNorm2d._cached_bboxs_len = len(WindowNorm2d._cached_bboxs)

self.num_features = num_features

self.mask_thres = mask_thres

self.input_size = input_size

self.mask_patch_size = mask_patch_size

self.eps = eps

self.device = device

# Mixup between window statistics and global statistics

self.mix = mix

self.alpha = torch.tensor(alpha, device=device)

self.beta_dist = torch.distributions.Beta(self.alpha, self.alpha)

# Grid mask (SimMIM-style) vs. random bounding box

self.grid = grid

self.affine = affine

if self.affine:

self.weight = nn.Parameter(torch.ones(num_features))

self.bias = nn.Parameter(torch.zeros(num_features))

else:

self.register_parameter("weight", None)

self.register_parameter("bias", None)

def forward(self, x):

if self.training:

N, C, H, W = x.size()

if self.grid:

# Lazy-initialise the fast grid mask generator on first call

if not hasattr(self, "_mask_generator"):

scale = round(self.input_size / H)

self._mask_generator = MaskGenerator(

input_size=self.input_size,

mask_patch_size=self.mask_patch_size,

model_patch_size=scale,

mask_ratio=self.mask_thres,

device=self.device,

)

masked_x, _ = self._mask_generator(x)

else:

if WindowNorm2d._cached_bboxs is None:

bbx1, bby1, bbx2, bby2 = cn_rand_bbox(

x.size(), beta=1, bbx_thres=self.mask_thres

)

else:

idx = np.random.randint(0, WindowNorm2d._cached_bboxs_len)

bbx1, bby1, bbx2, bby2 = WindowNorm2d._cached_bboxs[idx]

bbx1, bby1 = int(W * bbx1), int(H * bby1)

bbx2, bby2 = int(W * bbx2), int(H * bby2)

masked_x = x[:, :, bbx1:bbx2, bby1:bby2]

std, mean = calc_ins_std_mean(masked_x)

if self.mix and not self.grid:

lmda = self.beta_dist.sample((N, C, 1, 1))

global_std, global_mean = calc_ins_std_mean(x)

mean = mean * lmda + global_mean * (1 - lmda)

std = std * lmda + global_std * (1 - lmda)

else:

std, mean = calc_ins_std_mean(x)

normalized = (x - mean) / (std + self.eps)

if self.affine:

return self.weight[:, None, None] * normalized + self.bias[:, None, None]

return normalized

def __repr__(self):

return (

f"WindowNorm2d("

f"num_features={self.num_features}, "

f"mask_thres={self.mask_thres}, "

f"alpha={self.alpha.item():.4f}, "

f"mix={self.mix}, "

f"grid={self.grid}, "

f"input_size={self.input_size}, "

f"mask_patch_size={self.mask_patch_size}, "

f"eps={self.eps}, "

f"affine={self.affine}, "

f"cached={WindowNorm2d._cached_bboxs is not None}, "

f"device={self.device!r})"

)

# ------------------------------------------------------------------

# Model-conversion class methods

# ------------------------------------------------------------------

@classmethod

def convert_WIN_model(

cls,

module,

mask_thres=0.7,

alpha=0.1,

mix=True,

grid=False,

input_size=224,

mask_patch_size=32,

affine=False,

cached=False,

bboxs_path="./bboxs.npy",

device="cuda",

):

"""Recursively replace all ``BatchNorm2d`` layers with

mod = module

if isinstance(module, nn.BatchNorm2d):

mod = cls(

module.num_features,

mask_thres=mask_thres,

alpha=alpha,

mix=mix,

grid=grid,

input_size=input_size,

mask_patch_size=mask_patch_size,

affine=affine,

cached=cached,

bboxs_path=bboxs_path,

device=device,

)

for name, child in module.named_children():

mod.add_module(

name,

cls.convert_WIN_model(

child,

mask_thres=mask_thres,

alpha=alpha,

mix=mix,

grid=grid,

input_size=input_size,

mask_patch_size=mask_patch_size,

affine=affine,

cached=cached,

bboxs_path=bboxs_path,

device=device,

),

)

del module

return mod

@classmethod

def convert_IN_model(cls, module):

"""Recursively replace all ``BatchNorm2d`` layers with

mod = module

if isinstance(module, nn.BatchNorm2d):

mod = nn.InstanceNorm2d(module.num_features)

for name, child in module.named_children():

mod.add_module(name, cls.convert_IN_model(child))

del module

return mod

@classmethod

def convert_Identity_model(cls, module):

"""Recursively replace all ``BatchNorm2d`` layers with

mod = module

if isinstance(module, nn.BatchNorm2d):

mod = nn.Identity()

for name, child in module.named_children():

mod.add_module(name, cls.convert_Identity_model(child))

del module

return mod

@classmethod

def convert_GN_model(cls, module, num_groups=64):

"""Recursively replace all ``BatchNorm2d`` layers with

mod = module

if isinstance(module, nn.BatchNorm2d):

mod = nn.GroupNorm(num_groups, module.num_features)

for name, child in module.named_children():

mod.add_module(name, cls.convert_GN_model(child, num_groups=num_groups))

del module

"""Structured grid mask generator with a pre-computed mask pool.

def __init__(

self,

input_size=224,

mask_patch_size=32,

model_patch_size=2,

mask_ratio=0.6,

pool_size=512,

device="cpu",

):

assert input_size % mask_patch_size == 0, (

"input_size must be divisible by mask_patch_size"

)

assert mask_patch_size % model_patch_size == 0, (

"mask_patch_size must be divisible by model_patch_size"

)

self.rand_size = input_size // mask_patch_size

self.scale = mask_patch_size // model_patch_size

self.feat_size = self.rand_size * self.scale

self.token_count = self.rand_size ** 2

self.mask_count = int(np.ceil(self.token_count * mask_ratio))

self.keep_count = self.token_count - self.mask_count

self.pool_size = pool_size

self.device = device

# ── build mask pool ──────────────────────────────────────────────

# pool_flat: (P, token_count) True = masked

pool_flat = torch.zeros(pool_size, self.token_count, dtype=torch.bool)

for i in range(pool_size):

idx = torch.randperm(self.token_count)[: self.mask_count]

pool_flat[i, idx] = True

# Expand patch-level mask → feature-map resolution: (P, H_feat, W_feat)

pool_patch = pool_flat.view(pool_size, self.rand_size, self.rand_size)

pool_feat = (

pool_patch

.repeat_interleave(self.scale, dim=1)

.repeat_interleave(self.scale, dim=2)

)

# Pre-compute flat keep-indices: (P, K) where K = keep_count * scale²

keep_flat = ~pool_feat.view(pool_size, -1) # True = keep

keep_indices = [

keep_flat[i].nonzero(as_tuple=False).squeeze(1)

for i in range(pool_size)

]

# Stack requires all rows to have the same length — guaranteed because

# every mask has exactly mask_count masked patches × scale² pixels.

self.keep_indices = torch.stack(keep_indices).to(device) # (P, K)

self.pool_feat = pool_feat.to(device) # (P, H, W)

def __call__(self, x):

"""Apply a randomly selected pre-computed mask to feature map ``x``.

p = int(torch.randint(self.pool_size, (1,)).item())

mask = self.pool_feat[p] # (H, W)

N, C, H, W = x.shape

flat = x.reshape(N, C, H * W) # (N, C, H*W)

ki = self.keep_indices[p] # (K,)

ki_exp = ki.unsqueeze(0).unsqueeze(0).expand(N, C, -1)

masked_x = flat.gather(2, ki_exp) # (N, C, K)