Modifying YOLOv5 engine to work with custom classes

[SmartMilk]

Is there an easy way to setup the object tracker to track object classes outwith the COCO-128 dataset?

I have a trained YOLOv5 model with weights yolo5_custom.pt trained on a specific set of classes. However the YOLOv5 detection engine does not take a data or names input as an additional argument, and if I try modifying the coco128.yaml file contained within the data folder it throws an error when executing:

Input:
zero-shot-object-tracking/clip_object_tracker.py --source zero-shot-object-tracking/data/video/*******.mp4 --weights zero-shot-object-tracking/models/yolo5_custom.pt --detection-engine yolov5 --info

Output:

/bin/bash: /anaconda/envs/jupyter_env/lib/libtinfo.so.6: no version information available (required by /bin/bash)
Namespace(agnostic_nms=False, api_key=None, augment=False, cfg='yolov4.cfg', classes=None, confidence=0.4, detection_engine='yolov5', device='', exist_ok=False, img_size=640, info=True, max_cosine_distance=0.4, name='exp', names='zero-shot-object-tracking/coco-copy.names', nms_max_overlap=1.0, nn_budget=None, overlap=0.3, project='runs/detect', save_conf=False, save_txt=False, source='zero-shot-object-tracking/data/video/**********.mp4', thickness=3, update=False, url=None, view_img=False, weights=['zero-shot-object-tracking/models/model.pt'])
Traceback (most recent call last):
  File "zero-shot-object-tracking/clip_object_tracker.py", line 370, in <module>
    detect()
  File "zero-shot-object-tracking/clip_object_tracker.py", line 105, in detect
    yolov5_engine = Yolov5Engine(opt.weights, device, opt.classes, opt.confidence, opt.overlap, opt.agnostic_nms, opt.augment, half)
  File "/mnt/batch/tasks/shared/LS_root/mounts/clusters/compute-optimized-cpu001/code/Users/alex.jamieson/object_tracking_roboflow_method/zero-shot-object-tracking/utils/yolov5.py", line 6, in __init__
    self.model = attempt_load(weights, map_location=device)
  File "/mnt/batch/tasks/shared/LS_root/mounts/clusters/compute-optimized-cpu001/code/Users/alex.jamieson/object_tracking_roboflow_method/zero-shot-object-tracking/models/experimental.py", line 118, in attempt_load
    model.append(torch.load(w, map_location=map_location)['model'].float().fuse().eval())  # load FP32 model
KeyError: 'model'

Modified coco128.yaml file for reference:

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/coco128  # dataset root dir
train: images/train2017  # train images (relative to 'path') 128 images
val: images/train2017  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes
names:
  0: class1
  1: class2
  2: class3
  3: class4
  4: class5

# number of classes
nc: 5

# Download script/URL (optional)
download: https://ultralytics.com/assets/coco128.zip

edit: The root cause seems to come from my custom YOLOv5 model. Reverting the yaml config to its original state still throws the same error.