obichimav's picture
Upload 42 files
8e5d8c7 verified
# 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/
```