hf_README.md

🏠 Floorplan Segmentation Model Training

This repository contains the training code for a floorplan segmentation model that can identify walls, doors, windows, rooms, and background in architectural floorplans.

🎯 Model Architecture

• Type: Ultra Simple U-Net
• Input: RGB floorplan images (224x224)
• Output: 5-class segmentation (Background, Walls, Doors, Windows, Rooms)
• Parameters: ~258K

📊 Training Data

The model is trained on the Cubicasa5K dataset:

• Training: 4,200 images
• Validation: 400 images
• Test: 400 images

🚀 Quick Start

1. Setup Environment

pip install -r hf_requirements.txt

2. Prepare Data

1. Upload processed_data.zip to this repository
2. Extract the data: unzip processed_data.zip

3. Start Training

python hf_train.py

⚙️ Training Configuration

• Batch Size: 4
• Image Size: 224x224
• Epochs: 50
• Learning Rate: 1e-4
• Optimizer: Adam
• Loss: CrossEntropyLoss
• Scheduler: CosineAnnealingLR

📈 Expected Results

After training, you should see:

• Wall Coverage: 40-60% (vs previous 20.6%)
• Room Detection: Multiple rooms detected
• Door/Window Classification: Proper distinction from walls
• Overall Quality: Much better than previous attempts

💾 Model Outputs

• best_model.pth: Best trained model
• checkpoint_epoch_*.pth: Checkpoints every 10 epochs
• training_history.png: Training progress visualization

🔧 Usage

Load Trained Model

import torch
from hf_train import UltraSimpleModel

# Load model
model = UltraSimpleModel(n_channels=3, n_classes=5)
checkpoint = torch.load('best_model.pth', map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

Predict on New Image

import cv2
import torch

# Load and preprocess image
image = cv2.imread('floorplan.png')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image, (224, 224))
image_tensor = torch.from_numpy(image).float().permute(2, 0, 1) / 255.0
image_tensor = image_tensor.unsqueeze(0)

# Predict
with torch.no_grad():
    output = model(image_tensor)
    prediction = torch.argmax(output, dim=1).squeeze(0).numpy()

📊 Class Mapping

• 0: Background (Black)
• 1: Walls (Red)
• 2: Doors (Green)
• 3: Windows (Blue)
• 4: Rooms (Yellow)

🎯 Performance Metrics

• Loss: CrossEntropyLoss
• Validation: Every epoch
• Checkpointing: Every 10 epochs
• Best Model: Saved when validation loss improves

🔍 Troubleshooting

Common Issues

1. CUDA Out of Memory: Reduce batch size to 2
2. Data Not Found: Ensure processed_data.zip is uploaded
3. Slow Training: Check GPU availability

Performance Tips

• Use GPU for faster training
• Monitor GPU memory usage
• Clear cache periodically during training

📞 Support

If you encounter issues:

1. Check the training logs
2. Verify data format
3. Ensure all dependencies are installed

🏆 Results

This model should significantly improve upon the previous poor performance:

• Better wall detection
• Proper room segmentation
• Accurate door/window classification
• Overall higher quality results

---

Happy Training! 🚀