SOLO (Segmenting Objects by Locations) is an instance segmentation project that segments objects in an image using a unique approach where objects are localized and segmented in a unified process.
- datasets/: Contains scripts and data for handling datasets.
- images/: Includes sample images for inference and visualization.
- models/: Contains model architectures used for segmentation.
- trainer.py: The main training script for training the segmentation model.
- inference.py: Script for running inference on images using a pre-trained model.
- utils/: Utility functions for data preprocessing, post-processing, and visualization.
To run this project, ensure you have the following installed:
- Python 3.7 or higher
- PyTorch
- torchvision
- NumPy
- OpenCV
Install dependencies using:
pip install -r requirements.txt
- Training the model
To train the model, use the
trainer.pyscript:
python trainer.py
- Running Inference
To test the model on sample images, run:
python inference.py
The project outputs segmented images where each instance is highlighted distinctly.
Following are the loss curves for training and validation:
 |
 |
 |
| Total training loss | Focal loss (Category) for training data | Dice loss (Mask) for training data |
 |
 |
 |
| Total Validation loss | Focal loss (Category) for validation data | Dice loss (Mask) for validation data |
Below is an image showing a sample from the dataset used for training and validation:
Here are some example predictions made by the model:
 |
 |
 |
| Prediction 1 | Prediction 2 | Prediction 3 |
These predictions illustrate the model's ability to segment instances in the input images effectively, highlighting its performance on unseen data.
Customize the dataset and modify configurations as needed in the datasets/yolo_dataset.py script. The output from training and predictions can be visualized using the provided utility scripts.
For any questions or issues, please feel free to reach out or contribute to this project.