Training object detection model resultbook

Training object detection model resultbook#

此 colab 程式碼負責訓練模型。訓練好的模型可以在 Project_Nanshang or Predicting_mask_for_tifimage.ipynb 來建構地圖遮罩的 shp 檔。另外，本程式碼是從 Roboflow train-yolo-nas-on-custom-dataset.ipynb 而來，更多細節可以參考 Roboflow 的教程

⚡ Before you start#

Let’s make sure that we have access to GPU. We can use nvidia-smi command to do that. In case of any problems navigate to Edit -> Notebook settings -> Hardware accelerator, set it to GPU, and then click Save.

!nvidia-smi

NOTE: To make it easier for us to manage datasets, images and models we create a HOME constant.

import os
HOME = os.getcwd()
print(HOME)

Install YOLO-NAS#

!pip install -q super-gradients==3.1.3
!pip install -q roboflow
!pip install -q supervision

🚨 Restart the runtime

After installation is complete, you’ll need to restart the runtime after installation completes. Navigate to Runtime -> Restart runtime and confirm by clicking Yes when you see the popup.

After that, carry on with the notebook starting from the cell below

import torch

DEVICE = 'cuda' if torch.cuda.is_available() else "cpu"
MODEL_ARCH = 'yolo_nas_l'

from super_gradients.training import models

model = models.get(MODEL_ARCH, pretrained_weights="coco").to(DEVICE)

Login and Specify your dataset in Roboflow,#

In this project, we will also use roboflow’s tool to store our data. make sure you are already have a project on roboflow

請登入 roboflow 並寫上您的資料集名稱和版本

%cd {HOME}

import roboflow
from roboflow import Roboflow

roboflow.login()

rf = Roboflow()

Project_name = "nanshang_tomb"  # 請登入 roboflow 並寫上您的資料集名稱 和 版本
Version = 10

project = rf.workspace().project(Project_name)
dataset = project.version(Version).download("yolov5")

LOCATION = dataset.location
print("location:", LOCATION)
CLASSES = sorted(project.classes.keys())
print("classes:", CLASSES)

Fine-tune on custom dataset#

ckpt_root_dir - this is the directory where results from all your experiments will be saved

experiment_name - all checkpoints, logs, and tensorboards will be saved in a directory with the name you specify here.

settings#

MODEL_ARCH = 'yolo_nas_l'
BATCH_SIZE = 10
MAX_EPOCHS = 40
SCORE_THRSHOLD = 0.6
CHECKPOINT_DIR = f'{HOME}/checkpoints'
EXPERIMENT_NAME = project.name.lower().replace(" ", "_")

from super_gradients.training import Trainer

trainer = Trainer(experiment_name=EXPERIMENT_NAME, ckpt_root_dir=CHECKPOINT_DIR)

dataset_params = {
    'data_dir': LOCATION,
    'train_images_dir':'train/images',
    'train_labels_dir':'train/labels',
    'val_images_dir':'valid/images',
    'val_labels_dir':'valid/labels',
    'test_images_dir':'test/images',
    'test_labels_dir':'test/labels',
    'classes': CLASSES
}

from super_gradients.training.dataloaders.dataloaders import (
    coco_detection_yolo_format_train, coco_detection_yolo_format_val)

train_data = coco_detection_yolo_format_train(
    dataset_params={
        'data_dir': dataset_params['data_dir'],
        'images_dir': dataset_params['train_images_dir'],
        'labels_dir': dataset_params['train_labels_dir'],
        'classes': dataset_params['classes']
    },
    dataloader_params={
        'batch_size': BATCH_SIZE,
        'num_workers': 2
    }
)

val_data = coco_detection_yolo_format_val(
    dataset_params={
        'data_dir': dataset_params['data_dir'],
        'images_dir': dataset_params['val_images_dir'],
        'labels_dir': dataset_params['val_labels_dir'],
        'classes': dataset_params['classes']
    },
    dataloader_params={
        'batch_size': BATCH_SIZE,
        'num_workers': 2
    }
)

test_data = coco_detection_yolo_format_val(
    dataset_params={
        'data_dir': dataset_params['data_dir'],
        'images_dir': dataset_params['test_images_dir'],
        'labels_dir': dataset_params['test_labels_dir'],
        'classes': dataset_params['classes']
    },
    dataloader_params={
        'batch_size': BATCH_SIZE,
        'num_workers': 2
    }
)

Instantiate the model#

from super_gradients.training import models

model = models.get(
    MODEL_ARCH,
    num_classes=len(dataset_params['classes']),
    pretrained_weights="coco"
)

Define metrics and training parameters#

from super_gradients.training.losses import PPYoloELoss
from super_gradients.training.metrics import DetectionMetrics_050
from super_gradients.training.models.detection_models.pp_yolo_e import PPYoloEPostPredictionCallback

train_params = {
    'silent_mode': False,
    "average_best_models":True,
    "warmup_mode": "linear_epoch_step",
    "warmup_initial_lr": 1e-6,
    "lr_warmup_epochs": 3,
    "initial_lr": 5e-4,
    "lr_mode": "cosine",
    "cosine_final_lr_ratio": 0.1,
    "optimizer": "Adam",
    "optimizer_params": {"weight_decay": 0.0001},
    "zero_weight_decay_on_bias_and_bn": True,
    "ema": True,
    "ema_params": {"decay": 0.9, "decay_type": "threshold"},
    "max_epochs": MAX_EPOCHS,
    "mixed_precision": True,
    "loss": PPYoloELoss(
        use_static_assigner=False,
        num_classes=len(dataset_params['classes']),
        reg_max=16
    ),
    "valid_metrics_list": [
        DetectionMetrics_050(
            score_thres=SCORE_THRSHOLD,
            top_k_predictions=300,
            num_cls=len(dataset_params['classes']),
            normalize_targets=True,
            post_prediction_callback=PPYoloEPostPredictionCallback(
                score_threshold=0.01,
                nms_top_k=1000,
                max_predictions=300,
                nms_threshold=0.7
            )
        )
    ],
    "metric_to_watch": 'mAP@0.50'
}

Train the model#

trainer.train(
    model=model,
    training_params=train_params,
    train_loader=train_data,
    valid_loader=val_data
)

Analyze training metrics#

%load_ext tensorboard
%tensorboard --logdir {CHECKPOINT_DIR}/{EXPERIMENT_NAME}

Zip and download fine-tuned model#

# if you experience 'NotImplementedError: A UTF-8 locale is required. Got ANSI_X3.4-1968' error, run code below 👇

import locale
locale.getpreferredencoding = lambda: "UTF-8"

!zip -r yolo_nas.zip {CHECKPOINT_DIR}/{EXPERIMENT_NAME}

可以使用以下程式碼將資料轉移至自己的 google 雲端 ( address 要修改就是了 )

import locale
locale.getpreferredencoding = lambda: "UTF-8"
%cp "/content/checkpoints/nanshang_tomb/average_model.pth" "/content/drive/MyDrive/project_NanShang/resources/average_modelv10.pth"

記得將 checkpoints/nanshang_tomb 資料夾中的 average_model.pth 文件下載下來，此文件將在後續的物件偵測中被使用到。#

其本上 fine-tune 模型就到這裡就完成了。接下來可以到 Predicting_mask_for_tifimage.ipynb 來產生 shp 檔

Testing on testing dataset and Visualize them#

想要檢查 model 辨識能力的人請繼續閱讀，以下會在 testing dataset 中呈現預測的遮罩檔

Evaluate trained model#

trainer.test(
    model=model,
    test_loader=test_data,
    test_metrics_list=DetectionMetrics_050(
        score_thres=SCORE_THRSHOLD,
        top_k_predictions=300,
        num_cls=len(dataset_params['classes']),
        normalize_targets=True,
        post_prediction_callback=PPYoloEPostPredictionCallback(
            score_threshold=0.01,
            nms_top_k=1000,
            max_predictions=300,
            nms_threshold=0.7
        )
    )
)

Inference with trained model#

import supervision as sv

ds = sv.DetectionDataset.from_yolo(
    images_directory_path=f"{dataset.location}/test/images",
    annotations_directory_path=f"{dataset.location}/test/labels",
    data_yaml_path=f"{dataset.location}/data.yaml",
    force_masks=False
)

import supervision as sv

CONFIDENCE_TRESHOLD = 0.60

predictions = {}

for image_name, image in ds.images.items():
    result = list(model.predict(image, conf=CONFIDENCE_TRESHOLD))[0]
    detections = sv.Detections(
        xyxy=result.prediction.bboxes_xyxy,
        confidence=result.prediction.confidence,
        class_id=result.prediction.labels.astype(int)
    )
    predictions[image_name] = detections

Visualize inference results#

請調整 Seed 值以查看不同圖片的預測結果

import random

Seed = 14   # 可以調整 Seed 值以查看不同圖片的預測結果
random.seed(Seed)

import supervision as sv

MAX_IMAGE_COUNT = 5

n = min(MAX_IMAGE_COUNT, len(ds.images))

keys = list(ds.images.keys())
keys = random.sample(keys, n)

box_annotator = sv.BoxAnnotator()

images = []
titles = []

for key in keys:
    frame_with_annotations = box_annotator.annotate(
        scene=ds.images[key].copy(),
        detections=ds.annotations[key],
        skip_label=True
    )
    images.append(frame_with_annotations)
    titles.append('annotations')
    frame_with_predictions = box_annotator.annotate(
        scene=ds.images[key].copy(),
        detections=predictions[key],
        skip_label=True
    )
    images.append(frame_with_predictions)
    titles.append('predictions')

%matplotlib inline
sv.plot_images_grid(images=images, titles=titles, grid_size=(n, 2), size=(2 * 4, n * 4))

恭喜您完成訓練，如果您對模型呈現結果滿意，可以到下一個 workbook 來製作 shp 檔