# 🏠 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

```bash
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

```bash
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

```python
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

```python
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! 🚀**