Spaces:
Sleeping
Sleeping
File size: 3,288 Bytes
570db9a |
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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
# The model used in the paper Distribution Matching for Crowd Counting.
# Code adapted from https://github.com/cvlab-stonybrook/DM-Count/blob/master/models.py
from torch import nn, Tensor
import torch.nn.functional as F
from torch.hub import load_state_dict_from_url
from typing import Optional
from ..utils import make_vgg_layers, vgg_cfgs, vgg_urls
from ..utils import _init_weights
class VGG(nn.Module):
def __init__(
self,
features: nn.Module,
reduction: Optional[int] = None,
) -> None:
super().__init__()
self.features = features
self.reg_layer = nn.Sequential(
nn.Conv2d(512, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 128, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
)
self.reg_layer.apply(_init_weights)
# Remove the density layer, as the output from this model is not final and will be further processed.
# self.density_layer = nn.Sequential(nn.Conv2d(128, 1, 1), nn.ReLU())
self.encoder_reduction = 16
self.reduction = self.encoder_reduction if reduction is None else reduction
self.channels = 128
def forward(self, x: Tensor) -> Tensor:
x = self.features(x)
if self.encoder_reduction != self.reduction:
x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear")
x = self.reg_layer(x)
# x = self.density_layer(x)
return x
def _load_weights(model: VGG, url: str) -> VGG:
state_dict = load_state_dict_from_url(url)
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
print("Loading pre-trained weights")
if len(missing_keys) > 0:
print(f"Missing keys: {missing_keys}")
if len(unexpected_keys) > 0:
print(f"Unexpected keys: {unexpected_keys}")
return model
def vgg11(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["A"]), reduction=reduction)
return _load_weights(model, vgg_urls["vgg11"])
def vgg11_bn(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["A"], batch_norm=True), reduction=reduction)
return _load_weights(model, vgg_urls["vgg11_bn"])
def vgg13(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["B"]), reduction=reduction)
return _load_weights(model, vgg_urls["vgg13"])
def vgg13_bn(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["B"], batch_norm=True), reduction=reduction)
return _load_weights(model, vgg_urls["vgg13_bn"])
def vgg16(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["D"]), reduction=reduction)
return _load_weights(model, vgg_urls["vgg16"])
def vgg16_bn(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["D"], batch_norm=True), reduction=reduction)
return _load_weights(model, vgg_urls["vgg16_bn"])
def vgg19(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["E"]), reduction=reduction)
return _load_weights(model, vgg_urls["vgg19"])
def vgg19_bn(reduction: int = 8) -> VGG:
model = VGG(make_vgg_layers(vgg_cfgs["E"], batch_norm=True), reduction=reduction)
return _load_weights(model, vgg_urls["vgg19_bn"])
|