Spaces:
Runtime error
Runtime error
File size: 1,860 Bytes
f670afc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, check out LICENSE.md
from torch import nn
from torch.nn import functional as F
import torch.hub
class DeepLabV2(nn.Module):
def __init__(self, n_classes=182, image_size=512, use_dont_care=True):
super(DeepLabV2, self).__init__()
self.model = torch.hub.load(
"kazuto1011/deeplab-pytorch", "deeplabv2_resnet101",
pretrained=False, n_classes=182
)
state_dict = torch.hub.load_state_dict_from_url(
'https://github.com/kazuto1011/deeplab-pytorch/releases/download/'
'v1.0/deeplabv2_resnet101_msc-cocostuff164k-100000.pth',
map_location="cpu"
)
self.model.load_state_dict(state_dict)
self.image_size = image_size
# self.mean = torch.tensor([122.675, 116.669, 104.008], device="cuda")
self.mean = torch.tensor([104.008, 116.669, 122.675], device="cuda")
self.n_classes = n_classes
self.use_dont_care = use_dont_care
def forward(self, images, align_corners=True):
scale = self.image_size / max(images.shape[2:])
images = F.interpolate(
images, scale_factor=scale, mode='bilinear',
align_corners=align_corners
)
images = 255 * 0.5 * (images + 1) # (-1, 1) -> (0, 255)
images = images.flip(1) # RGB to BGR
images -= self.mean[None, :, None, None]
_, _, H, W = images.shape
logits = self.model(images)
logits = F.interpolate(
logits, size=(H, W), mode="bilinear",
align_corners=align_corners
)
probs = F.softmax(logits, dim=1)
pred = torch.argmax(probs, dim=1)
return pred
|