from segmentation_models_pytorch.encoders import encoders
from segmentation_models_pytorch import Unet
import torch

# Override pretrained settings for your weights
encoders["resnext50_32x4d"]["pretrained_settings"]["micronet"] = {
    "url": "https://huggingface.co/jstuckner/microscopy-resnext50-32x4d-micronet/resolve/main/resnext50_32x4d_micronet_weights.pth",
    "input_space": "RGB",
    "input_range": [0, 1],
    "mean": [0.485, 0.456, 0.406],
    "std": [0.229, 0.224, 0.225],
}

# Use as normal
model = Unet(
    encoder_name="resnext50_32x4d",
    encoder_weights="micronet",
    classes=1,
    activation=None,
)

# Test input
x = torch.randn(1, 3, 256, 256)
with torch.no_grad():
    y = model(x)
print("Output shape:", y.shape)