Views
No views yet
resnet34 backbone.1024x1024 pixels.segmentation_models_pytorch.1import torch
2import segmentation_models_pytorch as smp
3from huggingface_hub import hf_hub_download
4
5# Define the model architecture
6model = smp.Unet(
7 encoder_name="resnet34",
8 encoder_weights=None, # Weights are loaded from the saved state_dict
9 in_channels=3,
10 classes=1,
11 activation="sigmoid"
12)
13
14# Download and load state_dict
15model_path = hf_hub_download(repo_id="Frimpong1/pid-line-segmentation-unet-resnet34", filename="pytorch_model.bin")
16model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
17model.eval()
18
19# Example of using the model (assuming you have an image_tensor preprocessed)
20# image_tensor = ... # (C, H, W) float32 tensor, normalized
21# with torch.no_grad():
22# prediction = model(image_tensor.unsqueeze(0)) # Add batch dimension
23# mask = (prediction.squeeze(0).cpu().numpy() > 0.5).astype(float)