torchgeo/unet_resnet34_oam_rgb_tcd

Binary Segmentation U-Net ResNet-34 model for tree segmentation in OpenAerialMap high-resolution aerial imagery. The model is trained on a mix of CC-BY and CC-BY-NC licensed aerial imagery.

See for details: github.com/restor-foundation/tcd

Model weights extracted below:

import os
import hashlib

import torch
import segmentation_models_pytorch as smp


url = "https://huggingface.co/restor/tcd-unet-r34/resolve/main/model.pt"
state_dict = torch.hub.load_state_dict_from_url(url, weights_only=False, map_location="cpu")
model = smp.Unet(
    encoder_name="resnet34",
    encoder_weights=None,
    in_channels=3,
    classes=2,
)
model.load_state_dict(state_dict, strict=True)
filename = "unet_resnet34_oam_rgb_tcd.pth"
torch.save(model.state_dict(), filename)
md5 = hashlib.md5(open(filename, "rb").read()).hexdigest()[:8]
os.rename(filename, filename.replace(".pth", f"-{md5}.pth"))