phungpx/RMBG-1.4-wall-segmentation-cubicassa

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.