File size: 1,016 Bytes
8e5d8c7 |
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 |
# SemanticModel
Deep learning framework for semantic segmentation using PyTorch.
## Install
```bash
pip install -r requirements.txt
python setup.py install
```
## Usage
```python
from SemanticModel.model_core import SegmentationModel
from SemanticModel.prediction import PredictionPipeline
# Train
model = SegmentationModel(
classes=['background', 'object'],
architecture='unet',
encoder='timm-regnety_120'
)
trainer = ModelTrainer(
model_config=model,
root_dir='path/to/dataset',
epochs=40
)
model, metrics = trainer.train()
# Predict
predictor = PredictionPipeline(model)
predictor.predict_single_image('image.jpg')
predictor.predict_directory('image_dir/')
predictor.predict_raster('raster.tif')
# Load pretrained
model = SegmentationModel(
classes=['background', 'object'],
weights='path/to/best_model.pth'
)
```
## Data Structure
```
dataset/
├── train/
│ ├── Images/
│ └── Masks/
└── val/
├── Images/
└── Masks/
``` |