|
# 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/ |
|
``` |