Mercurial Hosting > traffic-intelligence
view scripts/dltrack.py @ 1222:69b531c7a061
added methods to reset trajectories and change object coordinates (including features)
author | Nicolas Saunier <nicolas.saunier@polymtl.ca> |
---|---|
date | Tue, 20 Jun 2023 15:42:19 -0400 |
parents | 5a207c838323 |
children | 759d76d6d20c |
line wrap: on
line source
#! /usr/bin/env python3 # from https://docs.ultralytics.com/modes/track/ import sys, argparse from trafficintelligence.moving import cocoUserTypeNames from ultralytics import YOLO parser = argparse.ArgumentParser(description='The program tracks objects following the ultralytics yolo executable.')#, epilog = 'Either the configuration filename or the other parameters (at least video and database filenames) need to be provided.') parser.add_argument('-i', dest = 'videoFilename', help = 'name of the video file (overrides the configuration file)') # detect model # tracker model parser.add_argument('--display', dest = 'display', help = 'show the results (careful with long videos, risk of running out of memory)', action = 'store_true') args = parser.parse_args() # Load a model model = YOLO('/home/nicolas/Research/Data/classification-models/yolov8x.pt') # seg yolov8x-seg.pt # seg could be used on cropped image... if can be loaded and kept in memory # model = YOLO('/home/nicolas/Research/Data/classification-models/yolo_nas_l.pt ') # AttributeError: 'YoloNAS_L' object has no attribute 'get' # Track with the model if args.display: results = model.track(source=args.videoFilename, tracker="/home/nicolas/Research/Data/classification-models/bytetrack.yaml", classes=list(cocoUserTypeNames.keys()), show=True) # , save_txt=True else: results = model.track(source=args.videoFilename, tracker="/home/nicolas/Research/Data/classification-models/bytetrack.yaml", classes=list(cocoUserTypeNames.keys()), stream=True) for result in results: for box in result.boxes: print(box.xyxy)