Views
No views yet
1pip install -r requirements.txt
2python setup.py install1from SemanticModel.model_core import SegmentationModel
2from SemanticModel.prediction import PredictionPipeline
3
4# Train
5model = SegmentationModel(
6 classes=['background', 'object'],
7 architecture='unet',
8 encoder='timm-regnety_120'
9)
10
11trainer = ModelTrainer(
12 model_config=model,
13 root_dir='path/to/dataset',
14 epochs=40
15)
16model, metrics = trainer.train()
17
18# Predict
19predictor = PredictionPipeline(model)
20predictor.predict_single_image('image.jpg')
21predictor.predict_directory('image_dir/')
22predictor.predict_raster('raster.tif')
23
24# Load pretrained
25model = SegmentationModel(
26 classes=['background', 'object'],
27 weights='path/to/best_model.pth'
28)dataset/
├── train/
│ ├── Images/
│ └── Masks/
└── val/
├── Images/
└── Masks/