--- 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