File size: 3,631 Bytes
6bd8c6d fb0ea68 6bd8c6d fb0ea68 6bd8c6d fb0ea68 6bd8c6d fb0ea68 6bd8c6d fb0ea68 6bd8c6d fb0ea68 | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | """
Self-contained UNetR50 for tree density estimation.
Dependencies: torch, segmentation-models-pytorch, huggingface_hub
Load a pretrained checkpoint:
from hub_model import UNetR50
model = UNetR50.from_pretrained("dgominski/TinyTrees", subfolder="ps") # Rwanda
model = UNetR50.from_pretrained("dgominski/TinyTrees", subfolder="gf") # China
model = UNetR50.from_pretrained("dgominski/TinyTrees", subfolder="spot") # France
Forward pass:
# x: (B, in_channels+1, H, W) — image bands + binary validity mask (last channel)
density_map = model(x) # (B, 1, H, W), trees/pixel
"""
from __future__ import annotations
import os
from pathlib import Path
from typing import Dict, Optional, Union
import torch
import torch.nn as nn
import segmentation_models_pytorch as smp
from huggingface_hub import PyTorchModelHubMixin, hf_hub_download
from huggingface_hub import constants as hf_constants
class UNetR50(
nn.Module,
PyTorchModelHubMixin,
library_name="treedensity",
repo_url="https://huggingface.co/dgominski/TinyTrees",
license="apache-2.0",
coders={},
):
"""ResNet-50 U-Net for per-pixel tree density regression.
Args:
in_channels: number of image bands (4 for PS/GF/SPOT). The model
expects in_channels+1 actual input channels: the image bands
concatenated with a binary validity mask as the last channel.
"""
def __init__(self, in_channels: int = 4, **kwargs):
super().__init__()
self.in_channels = in_channels
self.unet = smp.Unet(
encoder_name="resnet50",
encoder_weights=None,
in_channels=in_channels + 1,
classes=1,
activation=None,
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.unet(x)
def get_feats(self, x: torch.Tensor) -> torch.Tensor:
"""Return decoder feature map (used by P2PNet head)."""
return self.unet.decoder(self.unet.encoder(x))
@classmethod
def _from_pretrained(
cls,
*,
model_id: str,
revision: Optional[str],
cache_dir: Optional[Union[str, Path]],
force_download: bool,
proxies: Optional[Dict],
resume_download: Optional[bool],
local_files_only: bool,
token: Union[str, bool, None],
map_location: str = "cpu",
strict: bool = False,
subfolder: Optional[str] = None,
**model_kwargs,
):
model = cls(**model_kwargs)
def _dl(filename):
path_in_repo = f"{subfolder}/{filename}" if subfolder else filename
return hf_hub_download(
repo_id=model_id,
filename=path_in_repo,
revision=revision,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
resume_download=resume_download,
token=token,
local_files_only=local_files_only,
)
if os.path.isdir(model_id):
base = os.path.join(model_id, subfolder) if subfolder else model_id
sf = os.path.join(base, hf_constants.SAFETENSORS_SINGLE_FILE)
return cls._load_as_safetensor(model, sf, map_location, strict)
try:
return cls._load_as_safetensor(
model, _dl(hf_constants.SAFETENSORS_SINGLE_FILE), map_location, strict
)
except Exception:
return cls._load_as_pickle(
model, _dl(hf_constants.PYTORCH_WEIGHTS_NAME), map_location, strict
)
|