keystats/unetpp-resnest269e-survey

🧭 U-Net++ (ResNeSt269e) — Polygon Boundary Segmentation

This is a U-Net++ model trained for polygon boundary detection on survey plans using the timm-resnest269e encoder from segmentation_models_pytorch.

---

🧠 Model Details

• Architecture: U-Net++
• Encoder: timm-resnest269e
• Input size: 512 × 512
• Output: Binary mask
• Framework: PyTorch + segmentation_models_pytorch

⚡ Works on both CPU and GPU.

• CPU is slower but fully supported.
• GPU (cuda) significantly speeds up inference.

---

🚀 Quick Inference Example

#!pip install segmentation_models_pytorch huggingface_hub opencv-python matplotlib

import torch
import segmentation_models_pytorch as smp
from huggingface_hub import hf_hub_download
import json
import cv2
import numpy as np
import matplotlib.pyplot as plt

repo_id = "keystats/unetpp-resnest269e-survey"

config_path = hf_hub_download(repo_id, "config.json")
weights_path = hf_hub_download(repo_id, "pytorch_model.bin")

with open(config_path) as f:
    cfg = json.load(f)

model = smp.UnetPlusPlus(
    encoder_name=cfg["encoder_name"],
    encoder_weights=None,
    in_channels=cfg["in_channels"],
    classes=cfg["classes"],
    activation=cfg["activation"]
)

model.load_state_dict(torch.load(weights_path, map_location="cpu"))
model.eval()
print("✅ Model ready!")

# Path to your input image
image_path = "your survey image"

# Read image using OpenCV
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# Resize to model input size (e.g. 512x512)
IMAGE_SIZE = 512  # your training image size
image_resized = cv2.resize(image, (IMAGE_SIZE, IMAGE_SIZE))


# Normalize to [0, 1]
image_norm = image_resized / 255.0

# Convert to tensor: (H, W, C) -> (1, C, H, W)
tensor = torch.from_numpy(image_norm).permute(2, 0, 1).unsqueeze(0).float()

with torch.no_grad():
    output = model(tensor)  # shape: (1, 1, H, W)

# Convert prediction to numpy
mask = output.squeeze().cpu().numpy()

# Binarize the mask (0 or 1)
mask_binary = (mask > 0.5).astype(np.uint8) * 255

# Resize mask back to original image size if needed
mask_original_size = cv2.resize(mask_binary, (image.shape[1], image.shape[0]))

# Overlay or display
plt.figure(figsize=(10,5))
plt.subplot(1,2,1)
plt.title("Original Image")
plt.imshow(image)
plt.axis("off")

plt.subplot(1,2,2)
plt.title("Predicted Mask")
plt.imshow(mask_original_size, cmap="gray")
plt.axis("off")
plt.show()

# Optional: save the mask
cv2.imwrite("predicted_mask.png", mask_original_size)