2020-07-30 01:50:33 +00:00
|
|
|
import torch
|
|
|
|
|
from unet import UNet as _UNet
|
|
|
|
|
|
|
|
|
|
def unet_carvana(pretrained=False):
|
|
|
|
|
"""
|
|
|
|
|
UNet model trained on the Carvana dataset ( https://www.kaggle.com/c/carvana-image-masking-challenge/data ).
|
2021-08-19 09:14:08 +00:00
|
|
|
Set the scale to 0.5 (50%) when predicting.
|
2020-07-30 01:50:33 +00:00
|
|
|
"""
|
2021-08-19 09:14:08 +00:00
|
|
|
net = _UNet(n_channels=3, n_classes=2, bilinear=True)
|
2020-07-30 01:50:33 +00:00
|
|
|
if pretrained:
|
2021-08-19 09:14:08 +00:00
|
|
|
checkpoint = 'https://github.com/milesial/Pytorch-UNet/releases/download/v2.0/unet_carvana_scale0.5_epoch1.pth'
|
2020-07-30 01:50:33 +00:00
|
|
|
net.load_state_dict(torch.hub.load_state_dict_from_url(checkpoint, progress=True))
|
|
|
|
|
|
|
|
|
|
return net
|
|
|
|
|
|
