Merge pull request #38 from CornellDataScience/multithreading

Multithreading merged

src/processing/action.py

@@ -1,6 +1,7 @@
 from state import GameState, PlayerFrame, BallFrame, Box
 from args import DARGS
 
+
 class ActionRecognition:
     def __init__(self, state: GameState, args=DARGS) -> None:
         self.state = state
@@ -37,14 +38,13 @@ def pose_shot(self, player_frame: PlayerFrame, ANGLE_THRESHOLD):
                 curr += 0.075
         return curr
 
-
     def ball_shot(self, ball_frame: BallFrame, rim: Box):
         """
         Takes in ball data at a frame and returns the weighted sum of the ball
         features. Currently uses a simple threshold heuristic.
         """
         # checks whether ball y coord is above rim y coord
-        if ball_frame is None:
+        if ball_frame is None or rim is None:
             return 0
         ball_pos = ball_frame.box
         mid_box = ball_pos.center()
@@ -61,7 +61,6 @@ def ball_shot(self, ball_frame: BallFrame, rim: Box):
         # return y_weight + displacement_weight
         return y_weight + 0.1
 
-
     def shot_detect(self):
         """
         Iterates through each frame and computes the classification of shots for each player.
@@ -80,10 +79,12 @@ def shot_detect(self):
                     # shots_file.write(str(player_id))
                     # shots_file.write("\n")
                     for interval in self.state.possessions:
-                        if (interval.start <= frame.frameno <= interval.end and 
-                            interval not in self.state.shots):
+                        if (
+                            interval.start <= frame.frameno <= interval.end
+                            and interval not in self.state.shots
+                        ):
                             self.state.shots.append(interval)
-                    #shot_interval = (frame.frameno, player_id)
-                    #self.state.shots.append(shot_interval)
+                    # shot_interval = (frame.frameno, player_id)
+                    # self.state.shots.append(shot_interval)
 
         # shots_file.close()

src/strongsort/strong_sort/sort/track.py

@@ -66,8 +66,17 @@ class Track:
 
     """
 
-    def __init__(self, detection, track_id, class_id, conf, n_init, max_age, ema_alpha,
-                 feature=None):
+    def __init__(
+        self,
+        detection,
+        track_id,
+        class_id,
+        conf,
+        n_init,
+        max_age,
+        ema_alpha,
+        feature=None,
+    ):
         self.track_id = track_id
         self.class_id = int(class_id)
         self.hits = 1
@@ -117,13 +126,20 @@ def to_tlbr(self):
         ret[2:] = ret[:2] + ret[2:]
         return ret
 
-
-    def ECC(self, src, dst, warp_mode = cv2.MOTION_EUCLIDEAN, eps = 1e-5,
-        max_iter = 100, scale = 0.1, align = False):
+    def ECC(
+        self,
+        src,
+        dst,
+        warp_mode=cv2.MOTION_EUCLIDEAN,
+        eps=1e-5,
+        max_iter=100,
+        scale=0.1,
+        align=False,
+    ):
         """Compute the warp matrix from src to dst.
         Parameters
         ----------
-        src : ndarray 
+        src : ndarray
             An NxM matrix of source img(BGR or Gray), it must be the same format as dst.
         dst : ndarray
             An NxM matrix of target img(BGR or Gray).
@@ -151,10 +167,15 @@ def ECC(self, src, dst, warp_mode = cv2.MOTION_EUCLIDEAN, eps = 1e-5,
         """
 
         # skip if current and previous frame are not initialized (1st inference)
-        if (src.any() or dst.any() is None):
-            return None, None
-        # skip if current and previous fames are not the same size
-        elif (src.shape != dst.shape):
+        # if src or dst is None:
+        #     return None, None
+        try:
+            if src.any() or dst.any() is None:
+                return None, None
+            # skip if current and previous fames are not the same size
+            elif src.shape != dst.shape:
+                return None, None
+        except:
             return None, None
 
         # BGR2GRAY
@@ -167,16 +188,24 @@ def ECC(self, src, dst, warp_mode = cv2.MOTION_EUCLIDEAN, eps = 1e-5,
         if scale is not None:
             if isinstance(scale, float) or isinstance(scale, int):
                 if scale != 1:
-                    src_r = cv2.resize(src, (0, 0), fx = scale, fy = scale,interpolation =  cv2.INTER_LINEAR)
-                    dst_r = cv2.resize(dst, (0, 0), fx = scale, fy = scale,interpolation =  cv2.INTER_LINEAR)
+                    src_r = cv2.resize(
+                        src, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR
+                    )
+                    dst_r = cv2.resize(
+                        dst, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR
+                    )
                     scale = [scale, scale]
                 else:
                     src_r, dst_r = src, dst
                     scale = None
             else:
                 if scale[0] != src.shape[1] and scale[1] != src.shape[0]:
-                    src_r = cv2.resize(src, (scale[0], scale[1]), interpolation = cv2.INTER_LINEAR)
-                    dst_r = cv2.resize(dst, (scale[0], scale[1]), interpolation=cv2.INTER_LINEAR)
+                    src_r = cv2.resize(
+                        src, (scale[0], scale[1]), interpolation=cv2.INTER_LINEAR
+                    )
+                    dst_r = cv2.resize(
+                        dst, (scale[0], scale[1]), interpolation=cv2.INTER_LINEAR
+                    )
                     scale = [scale[0] / src.shape[1], scale[1] / src.shape[0]]
                 else:
                     src_r, dst_r = src, dst
@@ -185,20 +214,21 @@ def ECC(self, src, dst, warp_mode = cv2.MOTION_EUCLIDEAN, eps = 1e-5,
             src_r, dst_r = src, dst
 
         # Define 2x3 or 3x3 matrices and initialize the matrix to identity
-        if warp_mode == cv2.MOTION_HOMOGRAPHY :
+        if warp_mode == cv2.MOTION_HOMOGRAPHY:
             warp_matrix = np.eye(3, 3, dtype=np.float32)
-        else :
+        else:
             warp_matrix = np.eye(2, 3, dtype=np.float32)
 
         # Define termination criteria
         criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, max_iter, eps)
 
         # Run the ECC algorithm. The results are stored in warp_matrix.
         try:
-            (cc, warp_matrix) = cv2.findTransformECC (src_r, dst_r, warp_matrix, warp_mode, criteria, None, 1)
+            (cc, warp_matrix) = cv2.findTransformECC(
+                src_r, dst_r, warp_matrix, warp_mode, criteria, None, 1
+            )
         except cv2.error as e:
             return None, None
-
 
         if scale is not None:
             warp_matrix[0, 2] = warp_matrix[0, 2] / scale[0]
@@ -208,15 +238,18 @@ def ECC(self, src, dst, warp_mode = cv2.MOTION_EUCLIDEAN, eps = 1e-5,
             sz = src.shape
             if warp_mode == cv2.MOTION_HOMOGRAPHY:
                 # Use warpPerspective for Homography
-                src_aligned = cv2.warpPerspective(src, warp_matrix, (sz[1],sz[0]), flags=cv2.INTER_LINEAR)
-            else :
+                src_aligned = cv2.warpPerspective(
+                    src, warp_matrix, (sz[1], sz[0]), flags=cv2.INTER_LINEAR
+                )
+            else:
                 # Use warpAffine for Translation, Euclidean and Affine
-                src_aligned = cv2.warpAffine(src, warp_matrix, (sz[1],sz[0]), flags=cv2.INTER_LINEAR)
+                src_aligned = cv2.warpAffine(
+                    src, warp_matrix, (sz[1], sz[0]), flags=cv2.INTER_LINEAR
+                )
             return warp_matrix, src_aligned
         else:
             return warp_matrix, None
 
-
     def get_matrix(self, matrix):
         eye = np.eye(3)
         dist = np.linalg.norm(eye - matrix)
@@ -229,8 +262,8 @@ def camera_update(self, previous_frame, next_frame):
         warp_matrix, src_aligned = self.ECC(previous_frame, next_frame)
         if warp_matrix is None and src_aligned is None:
             return
-        [a,b] = warp_matrix
-        warp_matrix=np.array([a,b,[0,0,1]])
+        [a, b] = warp_matrix
+        warp_matrix = np.array([a, b, [0, 0, 1]])
         warp_matrix = warp_matrix.tolist()
         matrix = self.get_matrix(warp_matrix)
 
@@ -241,7 +274,6 @@ def camera_update(self, previous_frame, next_frame):
         cx, cy = x1_ + w / 2, y1_ + h / 2
         self.mean[:4] = [cx, cy, w / h, h]
 
-
     def increment_age(self):
         self.age += 1
         self.time_since_update += 1
@@ -270,11 +302,15 @@ def update(self, detection, class_id, conf):
         """
         self.conf = conf
         self.class_id = class_id.int()
-        self.mean, self.covariance = self.kf.update(self.mean, self.covariance, detection.to_xyah(), detection.confidence)
+        self.mean, self.covariance = self.kf.update(
+            self.mean, self.covariance, detection.to_xyah(), detection.confidence
+        )
 
         feature = detection.feature / np.linalg.norm(detection.feature)
 
-        smooth_feat = self.ema_alpha * self.features[-1] + (1 - self.ema_alpha) * feature
+        smooth_feat = (
+            self.ema_alpha * self.features[-1] + (1 - self.ema_alpha) * feature
+        )
         smooth_feat /= np.linalg.norm(smooth_feat)
         self.features = [smooth_feat]
 
@@ -284,16 +320,14 @@ def update(self, detection, class_id, conf):
             self.state = TrackState.Confirmed
 
     def mark_missed(self):
-        """Mark this track as missed (no association at the current time step).
-        """
+        """Mark this track as missed (no association at the current time step)."""
         if self.state == TrackState.Tentative:
             self.state = TrackState.Deleted
         elif self.time_since_update > self._max_age:
             self.state = TrackState.Deleted
 
     def is_tentative(self):
-        """Returns True if this track is tentative (unconfirmed).
-        """
+        """Returns True if this track is tentative (unconfirmed)."""
         return self.state == TrackState.Tentative
 
     def is_confirmed(self):

src/strongsort/yolov5/detect.py

@@ -50,6 +50,9 @@
 import torch
 import torch.backends.cudnn as cudnn
 
+import concurrent.futures
+import urllib.request
+
 FILE = Path(__file__).resolve()
 ROOT = FILE.parents[1]  # yolov5 strongsort root directory
 WEIGHTS = ROOT / "weights"
@@ -218,13 +221,30 @@ def run(
         with open(write_to, "w") as f:
             f.write("")
 
-    # Run tracking
-    model.warmup(imgsz=(1 if pt else nr_sources, 3, *imgsz))  # warmup
-    dt, seen = [0.0, 0.0, 0.0, 0.0], 0
-    curr_frames, prev_frames = [None] * nr_sources, [None] * nr_sources
-    for frame_idx, (path, im, im0s, vid_cap, s) in enumerate(dataset):
-        if not frame_idx % skips == 0:
-            continue
+    def runEverything(
+        frame_idx,
+        path,
+        im,
+        im0s,
+        vid_cap,
+        outputs,
+        s,
+        device,
+        half,
+        save_dir,
+        visualize,
+        augment,
+        model,
+        conf_thres,
+        iou_thres,
+        classes,
+        agnostic_nms,
+        max_det,
+        write_to,
+    ):
+        dt, seen = [0.0, 0.0, 0.0, 0.0], 0
+        # curr_frames, prev_frames = [None] * nr_sources, [None] * nr_sources
+
         t1 = time_sync()
         im = torch.from_numpy(im).to(device)
         im = im.half() if half else im.float()  # uint8 to fp16/32
@@ -274,6 +294,7 @@ def run(
                     )  # get folder name containing current img
                     save_path = str(save_dir / p.parent.name)  # im.jpg, vid.mp4, ...
             curr_frames[i] = im0
+            assert im0 is not None
 
             txt_path = str(save_dir / "tracks" / txt_file_name)  # im.txt
             if write_to is None:
@@ -523,6 +544,51 @@ def run(
                 vid_writer[i].write(im0)
 
             prev_frames[i] = curr_frames[i]
+            assert prev_frames[i] is not None
+        return dt, seen, save_path
+
+    # Run tracking
+    model.warmup(imgsz=(1 if pt else nr_sources, 3, *imgsz))  # warmup
+
+    # dt, seen = [0.0, 0.0, 0.0, 0.0], 0
+    curr_frames, prev_frames = [None] * nr_sources, [None] * nr_sources
+
+    # all = {executor.submit(runEverything, url, 60): url for url in URLS}
+    all = {}
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
+        for frame_idx, (path, im, im0s, vid_cap, s) in enumerate(dataset):
+            all[
+                executor.submit(
+                    runEverything,
+                    frame_idx,
+                    path,
+                    im,
+                    im0s,
+                    vid_cap,
+                    outputs,
+                    s,
+                    device,
+                    half,
+                    save_dir,
+                    visualize,
+                    augment,
+                    model,
+                    conf_thres,
+                    iou_thres,
+                    classes,
+                    agnostic_nms,
+                    max_det,
+                    write_to,
+                )
+            ] = [frame_idx, path, im, im0s, vid_cap, s]
+
+        dt, seen = [0.0, 0.0, 0.0, 0.0], 0
+        for future in concurrent.futures.as_completed(all):
+            [frame_idx, path, im, im0s, vid_cap, s] = all[future]
+            per_frame_dt, per_frame_seen, save_path = future.result()
+            dt += per_frame_dt
+            seen += per_frame_seen
 
     # Print results
     t = tuple(x / seen * 1e3 for x in dt)  # speeds per image