Log in to Google Colab, create a new document, change the running type to GPU
View allocated GPU (typically Tesla T4 16GB)
!nvidia-smi
Install the python library
%?pture !pip install super-gradients==3.1.0 !pip install imutils !pip install roboflow !pip install pytube --upgrade
call the API
from super_gradients.training import Trainer from super_gradients.training import dataloaders from super_gradients.training.dataloaders.dataloaders import coco_detection_yolo_format_train, coco_detection_yolo_format_val from IPython.display import clear_output from super_gradients.training.losses import PPYoloELosses from super_gradients.training.metrics import DetectionMetrics_050 from super_gradients.training.models.detection_models.pp_yolo_e import PPYoloEPostPredictionCallback from super_gradients. training import models
Set training experiment information
CHECKPOINT_DIR = 'checkpoints2' trainer = Trainer(experiment_name='ppe_yolonas_run2', ckpt_root_dir=CHECKPOINT_DIR)
Download the public dataset
!pip install roboflow from roboflow import Roboflow rf = Roboflow(api_key="4v0m0Q44ghy9SpeYcnuN") project = rf.workspace("objet-detect-yolov5").project("eep_detection-u9bbd") dataset = project.version(1).download("yolov5")
Set dataset information
dataset_params = { 'data_dir': '/content/EEP_Detection-1', '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': ['Protective Helmet', 'Shield', 'Jacket', 'Dust Mask', 'Eye Wear', 'Glove', 'Protective Boots'] }
Data loader
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': 16, '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': 16, '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': 16, 'num_workers': 2 } ) clear_output()
data conversion
train_data.dataset.transforms
train_data.dataset.dataset_params['transforms'][1]
train_data.dataset.dataset_params['transforms'][1]['DetectionRandomAffine']['degrees'] = 10.42
data visualization
train_data.dataset.plot()
Set up the model (eg YOLO-NAS S)
model = models.get('yolo_nas_s', num_classes=len(dataset_params['classes']), pretrained_weights="coco" )
Set training parameters
train_params = { # ENABLING SILENT MODE 'silent_mode': True, "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": 15, "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=0.1, 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": '[email protected]' }
Download example videos for post-training inference
!gdown "https://drive.google.com/uc?id=1crFwrpMF1OlaJ0ZCZjBNRo9llLEVR8VQ & amp;confirm=t" !gdown "https://drive.google.com/uc?id=1cTIBNQ1R_7JAOURVv9cJ6P935ym_IkZ0 & amp;confirm=t" !gdown "https://drive.google.com/uc?id=1256pNK0nQnEDT6FRLQAraTRkOY7BSprq & amp;confirm=t" !gdown "https://drive.google.com/uc?id=15D71z_g8uxZfXSx2ya3sy4n2-eg53meH & amp;confirm=t" !gdown "https://drive.google.com/uc?id=1iYW9ZAsYAaHkWZhFVwQh_ch41TMt30-Q & amp;confirm=t"
model training
trainer. train(model=model, training_params=train_params, train_loader = train_data, valid_loader=val_data)
save the best model
best_model = models.get('yolo_nas_s', num_classes=len(dataset_params['classes']), checkpoint_path="checkpoints2/ppe_yolonas_run2/ckpt_best.pth")
Test Results
trainer.test(model=best_model, test_loader=test_data, test_metrics_list=DetectionMetrics_050(score_thres=0.1, 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) ))
video reasoning
input_video_path = f"/content/demonew2.mp4" output_video_path = "detections.mp4" import torch device = 'cuda' if torch.cuda.is_available() else "cpu" best_model.to(device).predict(input_video_path).save(output_video_path)