Prediction on a video using YOLO-NAS trained on custom dataset crashing

[Dev-Devesu]

💡 Your Question

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

!pip install -q supervision

RESTART RUNTIME

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

DEVICE = 'cuda' if torch.cuda.is_available() else "cpu" MODEL_ARCH = 'yolo_nas_l' import super_gradients from super_gradients.training import models

model = models.get(MODEL_ARCH, pretrained_weights="coco").to(DEVICE) f"{HOME}/data" %cd {HOME} !mkdir {HOME}/data %cd {HOME}/data

import supervision as sv

detections = sv.Detections( xyxy=result.prediction.bboxes_xyxy, confidence=result.prediction.confidence, class_id=result.prediction.labels.astype(int) )

box_annotator = sv.BoxAnnotator()

labels = [ f"{result.class_names[classid]} {confidence:0.2f}" for , _, confidence, classid, in detections ]

annotated_frame = box_annotator.annotate( scene=image.copy(), detections=detections, labels=labels )

%matplotlib inline sv.plot_image(annotated_frame, (20, 20)) %cd {HOME}

import roboflow from roboflow import Roboflow

roboflow.login()

rf = Roboflow()

project = rf.workspace("roboflow-jvuqo").project("football-players-detection-2frwp") dataset = project.version(1).download("yolov5") LOCATION = dataset.location print("location:", LOCATION) CLASSES = sorted(project.classes.keys()) print("classes:", CLASSES) MODEL_ARCH = 'yolo_nas_l' BATCH_SIZE = 8 MAX_EPOCHS = 18 CHECKPOINT_DIR = f'{HOME}/checkpoints' EXPERIMENTNAME = 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 } ) train_data.dataset.transforms from super_gradients.training import models

model = models.get( MODEL_ARCH, num_classes=len(dataset_params['classes']), pretrained_weights="coco" ) 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=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": 'mAP@0.50' } %load_ext tensorboard %tensorboard --logdir {CHECKPOINT_DIR}/{EXPERIMENT_NAME} import locale locale.getpreferredencoding = lambda: "UTF-8" trainer.train( model=model, training_params=train_params, train_loader=train_data, valid_loader=val_data ) !zip -r yolo_nas.zip {CHECKPOINT_DIR}/{EXPERIMENT_NAME}

best_model = models.get( MODEL_ARCH, num_classes=len(dataset_params['classes']), checkpoint_path=f"{CHECKPOINT_DIR}/{EXPERIMENT_NAME}/average_model.pth" ).to(DEVICE) 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 ) ) ) !gdown "https://drive.google.com/uc?id=1PAZQliOpr2G7-jaHEDstzOu9CeoFeW9E&confirm=t" input_video_path = f"/content/08fd33_0.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)

My colab notebook is crashing as it exceeds the ram usage on the final line i.e when i am trying to use my model to predict on a demo video. Can i do something to avoid the crash?

Versions

No response

[Dev-Devesu]

My colab notebook is crashing as it exceeds the ram usage on the final line i.e when i am trying to use my model to predict on a demo video. Can i do something to avoid the crash?

[bit-scientist]

@Dev-Devesu just out of curiosity, how big is your 08fd33_0.mp4 in /content/?

[Dev-Devesu]

@Dev-Devesu just out of curiosity, how big is your 08fd33_0.mp4 in /content/?

~20mb. will it help if the file size is smaller?

[aymuos15]

Are you sure the colab account has not been used before (in the same day as your run in the issue) and its crashing because of excessive usage and not just this run?

[Dev-Devesu]

No, it's just from this run. Also it worked when i used a shorter video of smaller size.

[bit-scientist]

@Dev-Devesu Do you mean you resolved the above issue by using a shorter video. If so, go ahead and close the issus please.

[Dev-Devesu]

@bit-scientist I did solve the issue by using a shorter video but it would help a lot if I could make predictions of full length videos.

[bit-scientist]

Good to hear that, in terms of longer videos, maybe future enhancements will consider this too. I can leave it up to @Dev-Devesu