RMBG-1.4 — CubiCasa5K Binary Segmentation (wall)
Fine-tuned version of briaai/RMBG-1.4 on
phungpx/cubicassa5k-coco for binary
foreground/background segmentation of wall in architectural floor plans.
Evaluation — Binary metrics (validation split, threshold=0.5)
| Metric | Value |
|---|---|
| Pixel Accuracy | 0.9738 |
| IoU (Jaccard) | 0.7262 |
| Dice / F1 | 0.8414 |
| Precision | 0.8342 |
| Recall | 0.8486 |
Prediction samples
Columns: input floor-plan, ground-truth wall mask, and model prediction (per-image IoU shown above each prediction).
Training hyperparameters
| Hyperparameter | Value |
|---|---|
| base model | briaai/RMBG-1.4 |
| selected class | wall |
| image size | 1024 |
| batch size | 4 |
| epochs | 100 |
| learning rate | 1e-05 |
| weight decay | 0.05 |
| lr scheduler | cosine (warmup ratio 0.05) |
| loss | BCE + SSIM + IoU (weights {'bce': 1.0, 'ssim': 1.0, 'iou': 1.0}) |
| mixed precision | fp16 |
| seed | 42 |
Usage
import torch
import numpy as np
from PIL import Image
import torchvision.transforms as T
from transformers import AutoModelForImageSegmentation
model = AutoModelForImageSegmentation.from_pretrained("phungpx/RMBG-1.4-wall-segmentation-cubicassa", trust_remote_code=True)
model.eval()
transform = T.Compose([
T.Resize((1024, 1024)),
T.ToTensor(),
T.Normalize(mean=[0.5, 0.5, 0.5], std=[1.0, 1.0, 1.0]),
])
image = Image.open("floorplan.png").convert("RGB")
with torch.no_grad():
pred = model(pixel_values=transform(image).unsqueeze(0))
# Access prediction — adapt attribute based on model output:
mask = pred.pred_masks.squeeze().numpy() # (H, W) in [0,1]
binary_mask = (mask >= 0.4).astype(np.uint8)
Fine-tuned with the 🤗 Transformers custom Trainer subclass + BCE + SSIM + IoU loss.