File size: 3,295 Bytes
69f257e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 |
# 🏠 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! 🚀** |