Mercurial Hosting > traffic-intelligence
changeset 1233:d5695e0b59d9
saving results from ultralytics works
| author | Nicolas Saunier <nicolas.saunier@polymtl.ca> |
|---|---|
| date | Fri, 08 Sep 2023 17:09:12 -0400 |
| parents | 83ca1493d55c |
| children | dd969637381e |
| files | scripts/dltrack.py trafficintelligence/moving.py trafficintelligence/storage.py |
| diffstat | 3 files changed, 81 insertions(+), 28 deletions(-) [+] |
line wrap: on
line diff
--- a/scripts/dltrack.py Thu Sep 07 16:20:28 2023 -0400 +++ b/scripts/dltrack.py Fri Sep 08 17:09:12 2023 -0400 @@ -1,16 +1,18 @@ #! /usr/bin/env python3 # from https://docs.ultralytics.com/modes/track/ import sys, argparse - -from trafficintelligence import cvutils, moving, storage +from copy import copy from ultralytics import YOLO import cv2 +from trafficintelligence import cvutils, moving, storage, utils + 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') +#parser.add_argument('-f', dest = 'firstFrameNum', help = 'show the results (careful with long videos, risk of running out of memory)', action = 'store_true') args = parser.parse_args() # required functionality? @@ -53,45 +55,67 @@ # use 2 x bytetrack track buffer to remove objects from existing ones + +# check if one can go to specific frame https://docs.ultralytics.com/modes/track/#persisting-tracks-loop + # 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 +#results = model.track(source=args.videoFilename, tracker="/home/nicolas/Research/Data/classification-models/bytetrack.yaml", classes=list(moving.cocoTypeNames.keys()), show=True) # , save_txt=True if args.display: - results = model.track(source=args.videoFilename, tracker="/home/nicolas/Research/Data/classification-models/bytetrack.yaml", classes=list(moving.cocoTypeNames.keys()), show=True) # , save_txt=True -else: windowName = 'frame' cv2.namedWindow(windowName, cv2.WINDOW_NORMAL) - results = model.track(source=args.videoFilename, tracker="/home/nicolas/Research/Data/classification-models/bytetrack.yaml", classes=list(moving.cocoTypeNames.keys()), stream=True) - objects = [] - currentObjects = {} - featureNum = 0 - # create object with user type and list of 3 features (bottom ones and middle) + projection - for frameNum, result in enumerate(results): - print(frameNum, len(result.boxes)) - for box in result.boxes: +results = model.track(source=args.videoFilename, tracker="/home/nicolas/Research/Data/classification-models/bytetrack.yaml", classes=list(moving.cocoTypeNames.keys()), stream=True) +objects = [] +currentObjects = {} +featureNum = 0 +# create object with user type and list of 3 features (bottom ones and middle) + projection +for frameNum, result in enumerate(results): + print(frameNum, len(result.boxes)) + for box in result.boxes: + #print(box.cls, box.id, box.xyxy) + if box.id is not None: # None are objects with low confidence num = int(box.id) xyxy = box.xyxy[0].tolist() if num in currentObjects: currentObjects[num].timeInterval.last = frameNum - features = currentObjects[num].features - features[0].getPositions().addPositionXY(xyxy[0],xyxy[1]) - features[1].getPositions().addPositionXY(xyxy[2],xyxy[3]) + currentObjects[num].userTypes.append(moving.coco2Types[int(box.cls)]) + currentObjects[num].features[0].tmpPositions[frameNum] = moving.Point(xyxy[0],xyxy[1]) + currentObjects[num].features[1].tmpPositions[frameNum] = moving.Point(xyxy[2],xyxy[3]) + #features[0].getPositions().addPositionXY(xyxy[0],xyxy[1]) + #features[1].getPositions().addPositionXY(xyxy[2],xyxy[3]) else: - currentObjects[num] = moving.MovingObject(num, moving.TimeInterval(frameNum,frameNum), userType = moving.coco2Types[int(box.cls)]) - currentObjects[num].features = [moving.MovingObject(featureNum, moving.TimeInterval(frameNum, frameNum), moving.Trajectory([[xyxy[0]],[xyxy[1]]])), - moving.MovingObject(featureNum+1, moving.TimeInterval(frameNum, frameNum), moving.Trajectory([[xyxy[2]],[xyxy[3]]]))] + inter = moving.TimeInterval(frameNum,frameNum) + currentObjects[num] = moving.MovingObject(num, inter) + currentObjects[num].userTypes = [moving.coco2Types[int(box.cls)]] + currentObjects[num].features = [moving.MovingObject(featureNum), moving.MovingObject(featureNum+1)] currentObjects[num].featureNumbers = [featureNum, featureNum+1] + currentObjects[num].features[0].tmpPositions = {frameNum: moving.Point(xyxy[0],xyxy[1])} + currentObjects[num].features[1].tmpPositions = {frameNum: moving.Point(xyxy[2],xyxy[3])} featureNum += 2 - print(box.cls, box.xyxy) + if args.display: cvutils.cvImshow(windowName, result.plot()) # original image in orig_img key = cv2.waitKey() if cvutils.quitKey(key): break +# interpolate before saving +for num, obj in currentObjects.items(): + obj.setUserType(utils.mostCommon(obj.userTypes)) + obj.features[0].timeInterval = copy(obj.getTimeInterval()) + obj.features[1].timeInterval = copy(obj.getTimeInterval()) + if obj.length() != len(obj.features[0].tmpPositions): # interpolate + obj.features[0].positions = moving.Trajectory.fromPointDict(obj.features[0].tmpPositions) + obj.features[1].positions = moving.Trajectory.fromPointDict(obj.features[1].tmpPositions) + else: + obj.features[0].positions = moving.Trajectory.fromPointList(list(obj.features[0].tmpPositions.values())) + obj.features[1].positions = moving.Trajectory.fromPointList(list(obj.features[1].tmpPositions.values())) + storage.saveTrajectoriesToSqlite('test.sqlite', list(currentObjects.values()), 'object') # todo save bbox and mask to study localization / representation +# apply quality checks deviation and acceleration bounds?
--- a/trafficintelligence/moving.py Thu Sep 07 16:20:28 2023 -0400 +++ b/trafficintelligence/moving.py Fri Sep 08 17:09:12 2023 -0400 @@ -817,6 +817,34 @@ t.addPosition(p) return t + @staticmethod + def fromPointDict(points): + '''Points is a dict of points where keys are time instants + and there are (probably) missing positions''' + instants = sorted(list(points)) + # find all gaps + t1 = instants[0] + gap = False + gaps = [] + for t in range(instants[0], instants[-1]+1): + if t in instants: + if gap: + t2 = t + # store gap + gaps.append([t1, t2]) + gap = False + t1 = t + else: + t1 = t + else: + gap = True + # interpolate for gaps + for gap in gaps: + v = (points[gap[1]]-points[gap[0]]).divide(gap[1]-gap[0]) + for t in range(gap[0]+1, gap[1]): + points[t]=points[t-1]+v + return Trajectory.fromPointList([points[t] for t in range(instants[0], instants[-1]+1)]) + def __len__(self): return len(self.positions[0]) @@ -1376,13 +1404,13 @@ 'truck', 'automated'] -coco2Types = {0: 2, 1: 4, 2: 1, 5: 5, 7: 6} +coco2Types = {0: 2, 1: 4, 2: 1, 3: 3, 5: 5, 7: 6} cocoTypeNames = {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 5: 'bus', - 6: 'train', + #6: 'train', 7: 'truck'} userType2Num = utils.inverseEnumeration(userTypeNames)
--- a/trafficintelligence/storage.py Thu Sep 07 16:20:28 2023 -0400 +++ b/trafficintelligence/storage.py Fri Sep 08 17:09:12 2023 -0400 @@ -10,7 +10,7 @@ from numpy import log, min as npmin, max as npmax, round as npround, array, sum as npsum, loadtxt, floor as npfloor, ceil as npceil, linalg, int32, int64, reshape, dot, vstack, transpose, ones, zeros_like, pi, NaN from pandas import read_csv, merge, concat, DataFrame -from trafficintelligence import utils, moving, events, indicators, cvutils +from trafficintelligence import utils, moving, events, indicators from trafficintelligence.base import VideoFilenameAddable @@ -1307,6 +1307,7 @@ Ref: https://github.com/pratikac/kitti/blob/master/readme.tracking.txt''' from pykitti.utils import roty + from trafficintelligence.cvutils import cartesian2Homogeneous invR0 = linalg.inv(reshape(kittiCalibration['R_rect'], (3, 3))) transCam2Velo = transpose(kittiCalibration['Tr_cam_velo']) @@ -1371,16 +1372,16 @@ # corners3d = transpose(corners3d) # box3d_pts_3d_velo = calib.project_rect_to_velo(box3d_pts_3d) refCorners = transpose(dot(invR0, corners3d)) # 8x3 avoid double transpose np.transpose(pts_3d_rect))) in pts_3d_ref = self.project_rect_to_ref(pts_3d_rect) - homRefCorners = cvutils.cartesian2Homogeneous(refCorners) + homRefCorners = cartesian2Homogeneous(refCorners) veloCorners = dot(homRefCorners, transCam2Velo) # self.project_ref_to_velo(pts_3d_ref) # boxes3d_world_single, flag_imu, flag_world = calib_function.project_velo_to_world(bboxes_3d_velo=boxes3d_velo_single, Tr_matrix=transfer_matrix) #Tr_velo_to_imu = transfer_matrix['Tr_velo_to_imu'] #Tr_imu_to_world = transfer_matrix['Tr_imu_to_world'] - homVeloCorners = cvutils.cartesian2Homogeneous(veloCorners) + homVeloCorners = cartesian2Homogeneous(veloCorners) imuCorners = dot(kittiCalibration['Tr_velo_imu'], homVeloCorners.T).T # 8x3 - homImuCorners = cvutils.cartesian2Homogeneous(imuCorners) + homImuCorners = cartesian2Homogeneous(imuCorners) worldCorners = dot(Tr_imu_to_world, homImuCorners.T).T # 8x3 # take first 4 lines of corners, x,y,_ # x0, y0, _ = boxes3d[0] @@ -1415,12 +1416,12 @@ corners3d[1, :] = corners3d[1, :] corners3d[2, :] = corners3d[2, :] refCorners = transpose(dot(invR0, corners3d)) # 8x3 avoid double transpose np.transpose(pts_3d_rect))) in pts_3d_ref = self.project_rect_to_ref(pts_3d_rect) - homRefCorners = cvutils.cartesian2Homogeneous(refCorners) + homRefCorners = cartesian2Homogeneous(refCorners) veloCorners = dot(homRefCorners, transCam2Velo) - homVeloCorners = cvutils.cartesian2Homogeneous(veloCorners) + homVeloCorners = cartesian2Homogeneous(veloCorners) imuCorners = dot(kittiCalibration['Tr_velo_imu'], homVeloCorners.T).T # 8x3 - homImuCorners = cvutils.cartesian2Homogeneous(imuCorners) + homImuCorners = cartesian2Homogeneous(imuCorners) worldCorners = dot(Tr_imu_to_world, homImuCorners.T).T # 8x3 # take first 4 lines of corners, x,y,_ # x0, y0, _ = boxes3d[0]