add StrongSORT Tacker

feeder/trackers/deepocsort/cmc.py

@@ -0,0 +1,170 @@
+import pdb
+import pickle
+import os
+
+import cv2
+import numpy as np
+
+
+class CMCComputer:
+    def __init__(self, minimum_features=10, method="sparse"):
+        assert method in ["file", "sparse", "sift"]
+
+        os.makedirs("./cache", exist_ok=True)
+        self.cache_path = "./cache/affine_ocsort.pkl"
+        self.cache = {}
+        if os.path.exists(self.cache_path):
+            with open(self.cache_path, "rb") as fp:
+                self.cache = pickle.load(fp)
+        self.minimum_features = minimum_features
+        self.prev_img = None
+        self.prev_desc = None
+        self.sparse_flow_param = dict(
+            maxCorners=3000,
+            qualityLevel=0.01,
+            minDistance=1,
+            blockSize=3,
+            useHarrisDetector=False,
+            k=0.04,
+        )
+        self.file_computed = {}
+
+        self.comp_function = None
+        if method == "sparse":
+            self.comp_function = self._affine_sparse_flow
+        elif method == "sift":
+            self.comp_function = self._affine_sift
+        # Same BoT-SORT CMC arrays
+        elif method == "file":
+            self.comp_function = self._affine_file
+            self.file_affines = {}
+            # Maps from tag name to file name
+            self.file_names = {}
+
+            # All the ablation file names
+            for f_name in os.listdir("./cache/cmc_files/MOT17_ablation/"):
+                # The tag that'll be passed into compute_affine based on image name
+                tag = f_name.replace("GMC-", "").replace(".txt", "") + "-FRCNN"
+                f_name = os.path.join("./cache/cmc_files/MOT17_ablation/", f_name)
+                self.file_names[tag] = f_name
+            for f_name in os.listdir("./cache/cmc_files/MOT20_ablation/"):
+                tag = f_name.replace("GMC-", "").replace(".txt", "")
+                f_name = os.path.join("./cache/cmc_files/MOT20_ablation/", f_name)
+                self.file_names[tag] = f_name
+
+            # All the test file names
+            for f_name in os.listdir("./cache/cmc_files/MOTChallenge/"):
+                tag = f_name.replace("GMC-", "").replace(".txt", "")
+                if "MOT17" in tag:
+                    tag = tag + "-FRCNN"
+                # If it's an ablation one (not test) don't overwrite it
+                if tag in self.file_names:
+                    continue
+                f_name = os.path.join("./cache/cmc_files/MOTChallenge/", f_name)
+                self.file_names[tag] = f_name
+
+    def compute_affine(self, img, bbox, tag):
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        if tag in self.cache:
+            A = self.cache[tag]
+            return A
+        mask = np.ones_like(img, dtype=np.uint8)
+        if bbox.shape[0] > 0:
+            bbox = np.round(bbox).astype(np.int32)
+            bbox[bbox < 0] = 0
+            for bb in bbox:
+                mask[bb[1] : bb[3], bb[0] : bb[2]] = 0
+
+        A = self.comp_function(img, mask, tag)
+        self.cache[tag] = A
+
+        return A
+
+    def _load_file(self, name):
+        affines = []
+        with open(self.file_names[name], "r") as fp:
+            for line in fp:
+                tokens = [float(f) for f in line.split("\t")[1:7]]
+                A = np.eye(2, 3)
+                A[0, 0] = tokens[0]
+                A[0, 1] = tokens[1]
+                A[0, 2] = tokens[2]
+                A[1, 0] = tokens[3]
+                A[1, 1] = tokens[4]
+                A[1, 2] = tokens[5]
+                affines.append(A)
+        self.file_affines[name] = affines
+
+    def _affine_file(self, frame, mask, tag):
+        name, num = tag.split(":")
+        if name not in self.file_affines:
+            self._load_file(name)
+        if name not in self.file_affines:
+            raise RuntimeError("Error loading file affines for CMC.")
+
+        return self.file_affines[name][int(num) - 1]
+
+    def _affine_sift(self, frame, mask, tag):
+        A = np.eye(2, 3)
+        detector = cv2.SIFT_create()
+        kp, desc = detector.detectAndCompute(frame, mask)
+        if self.prev_desc is None:
+            self.prev_desc = [kp, desc]
+            return A
+        if desc.shape[0] < self.minimum_features or self.prev_desc[1].shape[0] < self.minimum_features:
+            return A
+
+        bf = cv2.BFMatcher(cv2.NORM_L2)
+        matches = bf.knnMatch(self.prev_desc[1], desc, k=2)
+        good = []
+        for m, n in matches:
+            if m.distance < 0.7 * n.distance:
+                good.append(m)
+
+        if len(good) > self.minimum_features:
+            src_pts = np.float32([self.prev_desc[0][m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
+            dst_pts = np.float32([kp[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
+            A, _ = cv2.estimateAffinePartial2D(src_pts, dst_pts, method=cv2.RANSAC)
+        else:
+            print("Warning: not enough matching points")
+        if A is None:
+            A = np.eye(2, 3)
+
+        self.prev_desc = [kp, desc]
+        return A
+
+    def _affine_sparse_flow(self, frame, mask, tag):
+        # Initialize
+        A = np.eye(2, 3)
+
+        # find the keypoints
+        keypoints = cv2.goodFeaturesToTrack(frame, mask=mask, **self.sparse_flow_param)
+
+        # Handle first frame
+        if self.prev_img is None:
+            self.prev_img = frame
+            self.prev_desc = keypoints
+            return A
+
+        matched_kp, status, err = cv2.calcOpticalFlowPyrLK(self.prev_img, frame, self.prev_desc, None)
+        matched_kp = matched_kp.reshape(-1, 2)
+        status = status.reshape(-1)
+        prev_points = self.prev_desc.reshape(-1, 2)
+        prev_points = prev_points[status]
+        curr_points = matched_kp[status]
+
+        # Find rigid matrix
+        if prev_points.shape[0] > self.minimum_features:
+            A, _ = cv2.estimateAffinePartial2D(prev_points, curr_points, method=cv2.RANSAC)
+        else:
+            print("Warning: not enough matching points")
+        if A is None:
+            A = np.eye(2, 3)
+
+        self.prev_img = frame
+        self.prev_desc = keypoints
+        return A
+
+    def dump_cache(self):
+        with open(self.cache_path, "wb") as fp:
+            pickle.dump(self.cache, fp)