calebkreed/bdd100k-unet-efficientnet-b0

Architecture (U-Net + EfficientNet-B0 via smp)
Dataset (BDD100K semantic segmentation, 19 classes)
Val mIoU (~0.54)

Training Details

• 25 Epochs
• Batches of 32
• Loss: CE + DICE
• Optimizer: Adam
• Learning Rate: 1e-4 with ReduceLROnPlateau

How to Use

pip install segmentation-models-pytorch albumentations opencv-python torch torchvision matplotlib

import torch
import cv2
import numpy as np
import albumentations as A
import segmentation_models_pytorch as smp

device = 'cuda' if torch.cuda.is_available() else 'cpu'

model_path = 'best_unet.pth'
img_path = 'your_image.jpg'

model = smp.Unet(
    encoder_name="efficientnet-b0",
    encoder_weights=None,
    in_channels=3,
    classes=19
).to(device)

state_dict = torch.load(model_path, map_location=device)
model.load_state_dict(state_dict)
model.eval()

class_names = ['road', 'sidewalk', 'building', 'wall', 'fence', 'pole',
               'traffic light', 'traffic sign', 'vegetation', 'terrain',
               'sky', 'person', 'rider', 'car', 'truck', 'bus',
               'train', 'motorcycle', 'bicycle']

img = cv2.imread(img_path)
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

pad = A.PadIfNeeded(min_height=736, min_width=1280, border_mode=cv2.BORDER_CONSTANT)
img_padded = pad(image=img_rgb)['image']

img_tensor = torch.from_numpy(img_padded).permute(2, 0, 1).float() / 255.0
img_tensor = img_tensor.unsqueeze(0).to(device)

with torch.no_grad():
    prediction = model(img_tensor)
    pred_mask = prediction.argmax(dim=1).squeeze(0).cpu().numpy()

# Visualization
cmap = plt.colormaps['tab20']
norm = plt.Normalize(vmin=0, vmax=18)

img_normalized = img_padded.astype(np.float32) / 255.0
pred_colored = cmap(norm(pred_mask))[..., :3]
overlay = (0.6 * img_normalized + 0.4 * pred_colored)

plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
plt.imshow(img_padded)
plt.title('Input Image')
plt.axis('off')

plt.subplot(1, 2, 2)
plt.imshow(overlay)
plt.title('Predicted Segmentation')
plt.axis('off')

plt.tight_layout()
plt.show()

Limitations

• Trained on BDD100K dashcam footage; performance may degrade on other camera angles or domains
• Struggles with thin/small structures (poles, fences, traffic signs)
• 2D only — no depth estimation
• Input images are padded to 736×1280; aspect ratios outside this may affect results