Yolov7

src/StrongSORT-YOLO/tracker.py

@@ -1,39 +1,51 @@
 from pathlib import Path
+import track_v7
 import sys
 import os
-import track_v5
 import pickle
 import multiprocessing as mp
 import threading as th
 import time
 
 FILE = Path(__file__).resolve()
-ROOT = FILE.parents[0]  # yolov5 strongsort root directory
-WEIGHTS = ROOT / 'weights'
-
-def track_person(res, source_mov:str, idx:int):
-    """ tracks persons in video and puts data in out_queue"""
-    out_array_pr, vid_path = track_v5.run(source = source_mov, classes= [1, 2],  yolo_weights = 
-                                          WEIGHTS / 'best.pt', save_vid=False, ret=True)
-
+ROOT = FILE.parents[0]
+WEIGHTS = ROOT / "weights"
+
+
+def track_person(res, source_mov: str, idx: int):
+    """tracks persons in video and puts data in out_queue"""
+    out_array_pr, vid_path = track_v7.run(
+        source=source_mov,
+        classes=[1, 2],
+        yolo_weights=WEIGHTS / "best.pt",
+        save_vid=False,
+        ret=True,
+    )
+
     res[idx] = (out_array_pr, vid_path)
 
     print("==============Put data from tracking person and rim============")
     return
 
-def track_basketball(res, source_mov:str, idx:int):
-    """ tracks basketball in video and puts data in out_queue"""
-    out_array_bb, bb_vid_path = track_v5.run(source = source_mov, 
-                                             yolo_weights = WEIGHTS / 'best_basketball.pt', 
-                                             save_vid=False, ret=True, skip_big=True)
-
+
+def track_basketball(res, source_mov: str, idx: int):
+    """tracks basketball in video and puts data in out_queue"""
+    out_array_bb, bb_vid_path = track_v7.run(
+        source=source_mov,
+        yolo_weights=WEIGHTS / "best_basketball.pt",
+        save_vid=False,
+        ret=True,
+        skip_big=True,
+    )
+
     res[idx] = (out_array_bb, bb_vid_path)
 
     print("==============Put data from tracking basketball============")
     return
 
-def get_data(source_mov:str):
-    """ returns dict as: {
+
+def get_data(source_mov: str):
+    """returns dict as: {
         'basketball_data': (basketball_bounding_boxes, video_path),
         'person_data': (persons_and_rim_bounding_boxes, vide_path)
         }.
@@ -42,7 +54,6 @@ def get_data(source_mov:str):
         source_mov (string path): path to video file
     """
 
-
     # ---------THREADING APPROACH---------------
 
     # res = [None] * 2
@@ -81,16 +92,20 @@ def get_data(source_mov:str):
 
     end = time.time()
 
-    print(f'=============time elapsed: {end-start}=================')
-
-    return {'basketball_data': (out_array_bb, bb_vid_path), 'person_data': (out_array_pr, vid_path)}
+    print(f"=============time elapsed: {end-start}=================")
+
+    return {
+        "basketball_data": (out_array_bb, bb_vid_path),
+        "person_data": (out_array_pr, vid_path),
+    }
 
 
-if __name__ == '__main__':
-    video_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', sys.argv[1])
+if __name__ == "__main__":
+    video_path = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "..", "..", sys.argv[1]
+    )
     output = get_data(video_path)
-    print(output['basketball_data'])
-    print('MODEL RUN DONE')
-    with open('../../tmp/test_output.pickle', 'wb') as f:
+    print(output["basketball_data"])
+    print("MODEL RUN DONE")
+    with open("../../tmp/test_output.pickle", "wb") as f:
         pickle.dump(output, f)
-

src/processing/court.py

@@ -567,7 +567,7 @@ def _max_pixel_overlap(self, im_src: np.ndarray, pts_src: list, pts_dst=None):
         """
         all_white = np.full_like(im_src, 255)
         max_overlap = self._apply_gray_homography(all_white, pts_src, pts_dst=pts_dst)
-        return np.count_nonzero(max_overlap > 100)
+        return np.count_nonzero(max_overlap > 100) + 1
 
     def _apply_bgr_homography(self, im_src: np.ndarray, pts_src: list):
         """