tracking.py 812 B

1234567891011121314
  1. from detection_helpers import *
  2. from tracking_helpers import *
  3. from bridge_wrapper import *
  4. from PIL import Image
  5. 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"
  6. detector.load_model('/root/helmet_det/yolov7-main/models/best.pt', img_size=2016) # pass the path to the trained weight file
  7. # Initialise class that binds detector and tracker in one class
  8. tracker = YOLOv7_DeepSORT(reID_model_path="./deep_sort/model_weights/mars-small128.pb", detector=detector)
  9. # output = None will not save the output video
  10. FILE = 'SD_Gate_Cam_1'
  11. 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)