| from segmentation_models_pytorch.encoders import encoders | |
| from segmentation_models_pytorch import Unet | |
| import torch | |
| # Override pretrained settings for your weights | |
| encoders["dpn98"]["pretrained_settings"]["micronet"] = { | |
| "url": "https://huggingface.co/jstuckner/microscopy-dpn98-micronet/resolve/main/dpn98_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="dpn98", | |
| 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) | |