Views
No views yet
timm-resnest269e encoder from segmentation_models_pytorch.cuda) significantly speeds up inference.1
2#!pip install segmentation_models_pytorch huggingface_hub opencv-python matplotlib
3
4import torch
5import segmentation_models_pytorch as smp
6from huggingface_hub import hf_hub_download
7import json
8import cv2
9import numpy as np
10import matplotlib.pyplot as plt
11
12repo_id = "keystats/unetpp-resnest269e-survey"
13
14config_path = hf_hub_download(repo_id, "config.json")
15weights_path = hf_hub_download(repo_id, "pytorch_model.bin")
16
17with open(config_path) as f:
18 cfg = json.load(f)
19
20model = smp.UnetPlusPlus(
21 encoder_name=cfg["encoder_name"],
22 encoder_weights=None,
23 in_channels=cfg["in_channels"],
24 classes=cfg["classes"],
25 activation=cfg["activation"]
26)
27
28model.load_state_dict(torch.load(weights_path, map_location="cpu"))
29model.eval()
30print("✅ Model ready!")
31
32# Path to your input image
33image_path = "your survey image"
34
35# Read image using OpenCV
36image = cv2.imread(image_path)
37image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
38
39# Resize to model input size (e.g. 512x512)
40IMAGE_SIZE = 512 # your training image size
41image_resized = cv2.resize(image, (IMAGE_SIZE, IMAGE_SIZE))
42
43
44# Normalize to [0, 1]
45image_norm = image_resized / 255.0
46
47# Convert to tensor: (H, W, C) -> (1, C, H, W)
48tensor = torch.from_numpy(image_norm).permute(2, 0, 1).unsqueeze(0).float()
49
50with torch.no_grad():
51 output = model(tensor) # shape: (1, 1, H, W)
52
53# Convert prediction to numpy
54mask = output.squeeze().cpu().numpy()
55
56# Binarize the mask (0 or 1)
57mask_binary = (mask > 0.5).astype(np.uint8) * 255
58
59# Resize mask back to original image size if needed
60mask_original_size = cv2.resize(mask_binary, (image.shape[1], image.shape[0]))
61
62# Overlay or display
63plt.figure(figsize=(10,5))
64plt.subplot(1,2,1)
65plt.title("Original Image")
66plt.imshow(image)
67plt.axis("off")
68
69plt.subplot(1,2,2)
70plt.title("Predicted Mask")
71plt.imshow(mask_original_size, cmap="gray")
72plt.axis("off")
73plt.show()
74
75# Optional: save the mask
76cv2.imwrite("predicted_mask.png", mask_original_size)