1234567891011121314 |
- from detection_helpers import *
- from tracking_helpers import *
- from bridge_wrapper import *
- from PIL import Image
- detector = Detector(classes = [0,1,2]) # it'll detect ONLY [helmet, people, head]. class = None means detect all classes. List info at: "data/coco.yaml"
- detector.load_model('/root/helmet_det/yolov7-main/models/best.pt', img_size=2016) # pass the path to the trained weight file
- # Initialise class that binds detector and tracker in one class
- tracker = YOLOv7_DeepSORT(reID_model_path="./deep_sort/model_weights/mars-small128.pb", detector=detector)
- # output = None will not save the output video
- FILE = 'SD_Gate_Cam_1'
- tracker.track_video(f"./IO_data/input/video/{FILE}.MP4", output=f"./IO_data/output/SD_Gate_Cam_2.mp4", show_live = False, skip_frames = None, count_objects = True, verbose=1)
|