pid-line-segmentation-unet-resnet34

P&ID Line Segmentation Model (U-Net with ResNet34 Backbone)

This model is a U-Net architecture with a ResNet34 encoder, fine-tuned for P&ID (Piping and Instrumentation Diagram) line segmentation.

Model Details

• Architecture: U-Net with resnet34 backbone.
• Input size: 1024x1024 pixels.
• Output: Binary mask (single class for lines).
• Loss Function: Combined Dice Loss and Focal Loss.

Performance (on Test Set)

• Test Loss: 0.5721
• Test Dice Score: 0.9413
• Test IoU Score: 0.8908

Usage

This model can be loaded using segmentation_models_pytorch.

import torch
import segmentation_models_pytorch as smp
from huggingface_hub import hf_hub_download

# Define the model architecture
model = smp.Unet(
    encoder_name="resnet34",
    encoder_weights=None, # Weights are loaded from the saved state_dict
    in_channels=3,
    classes=1,
    activation="sigmoid"
)

# Download and load state_dict
model_path = hf_hub_download(repo_id="Frimpong1/pid-line-segmentation-unet-resnet34", filename="pytorch_model.bin")
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
model.eval()

# Example of using the model (assuming you have an image_tensor preprocessed)
# image_tensor = ... # (C, H, W) float32 tensor, normalized
# with torch.no_grad():
#     prediction = model(image_tensor.unsqueeze(0)) # Add batch dimension
#     mask = (prediction.squeeze(0).cpu().numpy() > 0.5).astype(float)

Training Details

• Epochs Trained: 20
• Batch Size: 4
• Learning Rate: 0.0001
• Optimizer: Adam

License

(Specify your license here, e.g., MIT)