forked from milesial/Pytorch-UNet
-
Notifications
You must be signed in to change notification settings - Fork 0
/
hubconf.py
21 lines (17 loc) · 842 Bytes
/
hubconf.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import torch
from unet import UNet as _UNet
def unet_carvana(pretrained=False, scale=0.5):
"""
UNet model trained on the Carvana dataset ( https://www.kaggle.com/c/carvana-image-masking-challenge/data ).
Set the scale to 0.5 (50%) when predicting.
"""
net = _UNet(n_channels=3, n_classes=2, bilinear=False)
if pretrained:
if scale == 0.5:
checkpoint = 'https://github.com/milesial/Pytorch-UNet/releases/download/v3.0/unet_carvana_scale0.5_epoch2.pth'
elif scale == 1.0:
checkpoint = 'https://github.com/milesial/Pytorch-UNet/releases/download/v3.0/unet_carvana_scale1.0_epoch2.pth'
else:
raise RuntimeError('Only 0.5 and 1.0 scales are available')
net.load_state_dict(torch.hub.load_state_dict_from_url(checkpoint, progress=True))
return net