Yolov8 + Segnmentation

Project Website: https://github.com/WilliamSCUT/Yolov8-Segnmentation

This project mainly includes the implementation of a detection segmenter based on the yolov8-seg model.

Repo Structure

split_train_val.py Divide the training set, test set and validation set in proportion according to the existing labels.
text_to_yolo.py Convert xml format labels generated by labelImg to txt format.
auto_annotate.py Load the detection model and segmentation model, automatically annotate the specified data set, and generate labels with location and mask information.
train_v8.py Load the pre-trained model and train the corresponding detection and segmentation model on the specified data set.
test_v8.py Load the trained model for testing.

Installation

conda create -n yolo python=3.8
conda activate yolo
pip install labelImg
pip install ultralytics

Example Usages

To set up a new terminal, run:

conda activate yolo
cd <path to act repo>

Detection Data annotation

We use labelImg as the annotation tool. You need to open the images folder and change the save dir into your Annotation folder. To start it , run:

labelImg

Mask Auto Annotation

We use four_object_dection_model as the pretrained detection model, combining the segnmentation model sam_b.pt to achieve auto annotation. You only need to replace the detection model with your own model to achieve migration, and remember to change the model path in auto_annotate.py. To generated , run:

python auto_annotate.py

You can specify the label output path by changing output_dirs in this py file.

Segnmentation Model Training

To train the segnmentation model after get a large number of labeled images through auto annotation, run:

yolo segment train data=gear_classfication/gear_400.yaml model=yolov8s-seg.pt epochs=200 imgsz=640 name=seg_object

You can dynamically adjust parameters such as epochs according to the size and complexity of the data set, or modify the data path to your own data path.

Model Testing

To test the segnmentation model trained above, run:

yolo test_v8.py

You can specify the model path by changing model_path in this py file.

WilliamSCUT / Yolov8-Segnmentation