Mercurial Hosting > traffic-intelligence
diff scripts/classify-objects.py @ 1241:ab4c72b9475c
work in progress
author | Nicolas Saunier <nicolas.saunier@polymtl.ca> |
---|---|
date | Mon, 05 Feb 2024 17:06:01 -0500 |
parents | bb14f919d1cb |
children | 4cd8ace3552f |
line wrap: on
line diff
--- a/scripts/classify-objects.py Mon Feb 05 14:14:14 2024 -0500 +++ b/scripts/classify-objects.py Mon Feb 05 17:06:01 2024 -0500 @@ -5,6 +5,15 @@ import numpy as np import cv2 from scipy.stats import norm, lognorm +from pathlib import Path + +try: + from ultralytics import YOLO + ultralyticsAvailable = True +except ImportError: + #print('OpenCV library could not be loaded (video replay functions will not be available)') # TODO change to logging module + ultralyticsAvailable = False + from trafficintelligence import cvutils, moving, ml, storage, utils @@ -29,8 +38,15 @@ if speedAggregationFunc is None: sys.exit() -pedBikeCarSVM = ml.SVM_load(classifierParams.pedBikeCarSVMFilename) -bikeCarSVM = ml.SVM_load(classifierParams.bikeCarSVMFilename) +if ultralyticsAvailable and Path(classifierParams.dlFilename).is_file(): # use Yolo + pedBikeCarSVM = None + bikeCarSVM = None + yolo = YOLO(classifierParams.dlFilename, task='detect') + useYolo = True +else: + useYolo = False + pedBikeCarSVM = ml.SVM_load(classifierParams.pedBikeCarSVMFilename) + bikeCarSVM = ml.SVM_load(classifierParams.bikeCarSVMFilename) # log logistic for ped and bik otherwise ((pedBeta/pedAlfa)*((sMean/pedAlfa)**(pedBeta-1)))/((1+(sMean/pedAlfa)**pedBeta)**2.) carNorm = norm(classifierParams.meanVehicleSpeed, classifierParams.stdVehicleSpeed) @@ -61,8 +77,6 @@ objects = storage.loadTrajectoriesFromSqlite(databaseFilename, 'object', args.nObjects, withFeatures = True) timeInterval = moving.TimeInterval.unionIntervals([obj.getTimeInterval() for obj in objects]) -if args.startFrame0: - timeInterval.first = 0 capture = cv2.VideoCapture(videoFilename) width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) @@ -81,8 +95,7 @@ if capture.isOpened(): ret = True frameNum = timeInterval.first - if not args.startFrame0: - capture.set(cv2.CAP_PROP_POS_FRAMES, frameNum) + capture.set(cv2.CAP_PROP_POS_FRAMES, frameNum) lastFrameNum = timeInterval.last while ret and frameNum <= lastFrameNum: @@ -91,7 +104,10 @@ if frameNum%50 == 0: print('frame number: {}'.format(frameNum)) #if undistort: - # img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR) + # img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR) + if useYolo: + results = yolo.predict(img, classes=list(moving.cocoTypeNames.keys()), verbose=False) + for obj in objects[:]: if obj.getFirstInstant() <= frameNum: # if images are skipped obj.initClassifyUserTypeHoGSVM(speedAggregationFunc, pedBikeCarSVM, bikeCarSVM, classifierParams.maxPedestrianSpeed, classifierParams.maxCyclistSpeed, classifierParams.nFramesIgnoreAtEnds, invHomography, intrinsicCameraMatrix, distortionCoefficients) @@ -99,12 +115,18 @@ objects.remove(obj) for obj in currentObjects[:]: - if obj.getLastInstant() <= frameNum: # if images are skipped + if obj.getLastInstant() <= frameNum: obj.classifyUserTypeHoGSVM(minSpeedEquiprobable = classifierParams.minSpeedEquiprobable, speedProbabilities = speedProbabilities, maxPercentUnknown = classifierParams.maxPercentUnknown) pastObjects.append(obj) currentObjects.remove(obj) else: - obj.classifyUserTypeHoGSVMAtInstant(img, frameNum, width, height, classifierParams.percentIncreaseCrop, classifierParams.percentIncreaseCrop, classifierParams.minNPixels, classifierParams.hogRescaleSize, classifierParams.hogNOrientations, classifierParams.hogNPixelsPerCell, classifierParams.hogNCellsPerBlock, classifierParams.hogBlockNorm) + if useYolo: + # if one feature falls in bike, it's a bike + # could one count all hits in various objects, or one takes majority at the instant? + # obj.classifyUserTypeYoloAtInstant(img, frameNum, width, height, classifierParams.percentIncreaseCrop, classifierParams.percentIncreaseCrop, results[0].boxes) + pass + else: + obj.classifyUserTypeHoGSVMAtInstant(img, frameNum, width, height, classifierParams.percentIncreaseCrop, classifierParams.percentIncreaseCrop, classifierParams.minNPixels, classifierParams.hogRescaleSize, classifierParams.hogNOrientations, classifierParams.hogNPixelsPerCell, classifierParams.hogNCellsPerBlock, classifierParams.hogBlockNorm) if args.verbose: print('obj {}@{}: {}'.format(obj.getNum(), frameNum, moving.userTypeNames[obj.userTypes[frameNum]])) frameNum += 1