supervisely 6.73.417__py3-none-any.whl → 6.73.419__py3-none-any.whl

supervisely/nn/tracker/botsort/osnet_reid/osnet_reid_interface.py

@@ -0,0 +1,88 @@
+from pathlib import Path
+import cv2
+import numpy as np
+from .osnet import osnet_x1_0
+from collections import OrderedDict
+from supervisely import logger
+
+try:
+    # pylint: disable=import-error
+    import torch
+    from torch.nn import functional as F
+except ImportError:
+    logger.warning("torch is not installed, OSNet re-ID cannot be used.")
+
+
+class OsnetReIDModel:
+    def __init__(self, weights_path: Path = None, device: torch.device = torch.device("cpu"), half: bool = False):
+        self.device = device
+        self.half = half
+        self.input_shape = (256, 128)
+        if weights_path is None:
+            self.model = osnet_x1_0(num_classes=1000, loss='softmax', pretrained=True, use_gpu=device)
+        else:
+            self.model = osnet_x1_0(num_classes=1000, loss='softmax', pretrained=False, use_gpu=device)
+            self.load_pretrained_weights(weights_path)
+        self.model.to(self.device).eval()
+        if self.half:
+            self.model.half()
+
+    def load_pretrained_weights(self, weight_path: Path):
+        checkpoint = torch.load(weight_path, map_location=self.device)
+        state_dict = checkpoint.get("state_dict", checkpoint)
+        model_dict = self.model.state_dict()
+
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            key = k[7:] if k.startswith("module.") else k
+            if key in model_dict and model_dict[key].size() == v.size():
+                new_state_dict[key] = v
+
+        model_dict.update(new_state_dict)
+        self.model.load_state_dict(model_dict)
+
+    def get_features(self, xyxys, img: np.ndarray):
+        if xyxys.size == 0:
+            return np.empty((0, 512))
+
+        crops = self._get_crops(xyxys, img)
+        with torch.no_grad():
+            features = self.model(crops)
+            features = F.normalize(features, dim=1).cpu().numpy()
+
+        return features
+
+    def _get_crops(self, xyxys, img):
+        h, w = img.shape[:2]
+        crops = []
+
+        for box in xyxys:
+            x1, y1, x2, y2 = box.round().astype(int)
+            x1, y1, x2, y2 = max(0, x1), max(0, y1), min(w, x2), min(h, y2)
+            crop = cv2.resize(img[y1:y2, x1:x2], self.input_shape[::-1])
+            crop = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
+            crop = crop.astype(np.float32) / 255.0
+            crop = (crop - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
+            crop = torch.from_numpy(crop).permute(2, 0, 1)
+            crops.append(crop)
+
+        batch = torch.stack(crops).to(self.device, dtype=torch.float16 if self.half else torch.float32)
+        return batch
+
+
+
+class OsnetReIDInterface:
+    def __init__(self, weights: Path, device: str = "cpu", fp16: bool = False):
+        self.device = torch.device(device)
+        self.fp16 = fp16
+        self.model = OsnetReIDModel(weights, self.device, half=fp16)
+
+    def inference(self, image: np.ndarray, detections: np.ndarray) -> np.ndarray:
+        if detections is None or np.size(detections) == 0:
+            return np.zeros((0, 512), dtype=np.float32)  # пустой набор фичей
+
+        xyxys = detections[:, 0:4]  # left, top, right, bottom
+        features = self.model.get_features(xyxys, image)
+        return features
+
+ 

supervisely/nn/tracker/{bot_sort → botsort/tracker}/basetrack.py

@@ -1,6 +1,5 @@
-from collections import OrderedDict
-
 import numpy as np
+from collections import OrderedDict
 
 
 class TrackState(object):

supervisely/nn/tracker/{utils → botsort/tracker}/gmc.py

@@ -1,60 +1,59 @@
-import copy
-import time
-
 import cv2
 import matplotlib.pyplot as plt
 import numpy as np
+import copy
+import time
 
 
 class GMC:
-    def __init__(self, method="sparseOptFlow", downscale=2, gmc_file: str = None):
+    def __init__(self, method='sparseOptFlow', downscale=2, verbose=None):
         super(GMC, self).__init__()
 
         self.method = method
         self.downscale = max(1, int(downscale))
 
-        if self.method == "orb":
+        if self.method == 'orb':
             self.detector = cv2.FastFeatureDetector_create(20)
             self.extractor = cv2.ORB_create()
             self.matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
 
-        elif self.method == "sift":
-            self.detector = cv2.SIFT_create(
-                nOctaveLayers=3, contrastThreshold=0.02, edgeThreshold=20
-            )
-            self.extractor = cv2.SIFT_create(
-                nOctaveLayers=3, contrastThreshold=0.02, edgeThreshold=20
-            )
+        elif self.method == 'sift':
+            self.detector = cv2.SIFT_create(nOctaveLayers=3, contrastThreshold=0.02, edgeThreshold=20)
+            self.extractor = cv2.SIFT_create(nOctaveLayers=3, contrastThreshold=0.02, edgeThreshold=20)
             self.matcher = cv2.BFMatcher(cv2.NORM_L2)
 
-        elif self.method == "ecc":
+        elif self.method == 'ecc':
             number_of_iterations = 5000
             termination_eps = 1e-6
             self.warp_mode = cv2.MOTION_EUCLIDEAN
-            self.criteria = (
-                cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT,
-                number_of_iterations,
-                termination_eps,
-            )
-
-        elif self.method == "sparseOptFlow":
-            self.feature_params = dict(
-                maxCorners=1000,
-                qualityLevel=0.01,
-                minDistance=1,
-                blockSize=3,
-                useHarrisDetector=False,
-                k=0.04,
-            )
+            self.criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, number_of_iterations, termination_eps)
+
+        elif self.method == 'sparseOptFlow':
+            self.feature_params = dict(maxCorners=1000, qualityLevel=0.01, minDistance=1, blockSize=3,
+                                       useHarrisDetector=False, k=0.04)
             # self.gmc_file = open('GMC_results.txt', 'w')
 
-        elif self.method == "file" or self.method == "files":
-            self.gmcFile = open(gmc_file, "r")
-            if self.gmcFile is None:
-                raise ValueError("Error: Unable to open GMC file:" + gmc_file)
+        elif self.method == 'file' or self.method == 'files':
+            seqName = verbose[0]
+            ablation = verbose[1]
+            if ablation:
+                filePath = r'tracker/GMC_files/MOT17_ablation'
+            else:
+                filePath = r'tracker/GMC_files/MOTChallenge'
+
+            if '-FRCNN' in seqName:
+                seqName = seqName[:-6]
+            elif '-DPM' in seqName:
+                seqName = seqName[:-4]
+            elif '-SDP' in seqName:
+                seqName = seqName[:-4]
+
+            self.gmcFile = open(filePath + "/GMC-" + seqName + ".txt", 'r')
 
-        elif self.method is None or self.method in ("None", "none"):
-            self.method = "none"
+            if self.gmcFile is None:
+                raise ValueError("Error: Unable to open GMC file in directory:" + filePath)
+        elif self.method == 'none' or self.method == 'None':
+            self.method = 'none'
         else:
             raise ValueError("Error: Unknown CMC method:" + method)
 
@@ -65,15 +64,15 @@ class GMC:
         self.initializedFirstFrame = False
 
     def apply(self, raw_frame, detections=None):
-        if self.method == "orb" or self.method == "sift":
+        if self.method == 'orb' or self.method == 'sift':
             return self.applyFeaures(raw_frame, detections)
-        elif self.method == "ecc":
+        elif self.method == 'ecc':
             return self.applyEcc(raw_frame, detections)
-        elif self.method == "sparseOptFlow":
+        elif self.method == 'sparseOptFlow':
             return self.applySparseOptFlow(raw_frame, detections)
-        elif self.method == "file":
+        elif self.method == 'file':
             return self.applyFile(raw_frame, detections)
-        elif self.method == "none":
+        elif self.method == 'none':
             return np.eye(2, 3)
         else:
             return np.eye(2, 3)
@@ -105,11 +104,9 @@ class GMC:
         # Run the ECC algorithm. The results are stored in warp_matrix.
         # (cc, H) = cv2.findTransformECC(self.prevFrame, frame, H, self.warp_mode, self.criteria)
         try:
-            (cc, H) = cv2.findTransformECC(
-                self.prevFrame, frame, H, self.warp_mode, self.criteria, None, 1
-            )
+            (cc, H) = cv2.findTransformECC(self.prevFrame, frame, H, self.warp_mode, self.criteria, None, 1)
         except:
-            print("Warning: find transform failed. Set warp as identity")
+            print('Warning: find transform failed. Set warp as identity')
 
         return H
 
@@ -130,11 +127,11 @@ class GMC:
         # find the keypoints
         mask = np.zeros_like(frame)
         # mask[int(0.05 * height): int(0.95 * height), int(0.05 * width): int(0.95 * width)] = 255
-        mask[int(0.02 * height) : int(0.98 * height), int(0.02 * width) : int(0.98 * width)] = 255
+        mask[int(0.02 * height): int(0.98 * height), int(0.02 * width): int(0.98 * width)] = 255
         if detections is not None:
             for det in detections:
                 tlbr = (det[:4] / self.downscale).astype(np.int_)
-                mask[tlbr[1] : tlbr[3], tlbr[0] : tlbr[2]] = 0
+                mask[tlbr[1]:tlbr[3], tlbr[0]:tlbr[2]] = 0
 
         keypoints = self.detector.detect(frame, mask)
 
@@ -176,14 +173,11 @@ class GMC:
                 prevKeyPointLocation = self.prevKeyPoints[m.queryIdx].pt
                 currKeyPointLocation = keypoints[m.trainIdx].pt
 
-                spatialDistance = (
-                    prevKeyPointLocation[0] - currKeyPointLocation[0],
-                    prevKeyPointLocation[1] - currKeyPointLocation[1],
-                )
+                spatialDistance = (prevKeyPointLocation[0] - currKeyPointLocation[0],
+                                   prevKeyPointLocation[1] - currKeyPointLocation[1])
 
-                if (np.abs(spatialDistance[0]) < maxSpatialDistance[0]) and (
-                    np.abs(spatialDistance[1]) < maxSpatialDistance[1]
-                ):
+                if (np.abs(spatialDistance[0]) < maxSpatialDistance[0]) and \
+                        (np.abs(spatialDistance[1]) < maxSpatialDistance[1]):
                     spatialDistances.append(spatialDistance)
                     matches.append(m)
 
@@ -233,7 +227,7 @@ class GMC:
                 H[0, 2] *= self.downscale
                 H[1, 2] *= self.downscale
         else:
-            print("Warning: not enough matching points")
+            print('Warning: not enough matching points')
 
         # Store to next iteration
         self.prevFrame = frame.copy()
@@ -271,9 +265,7 @@ class GMC:
             return H
 
         # find correspondences
-        matchedKeypoints, status, err = cv2.calcOpticalFlowPyrLK(
-            self.prevFrame, frame, self.prevKeyPoints, None
-        )
+        matchedKeypoints, status, err = cv2.calcOpticalFlowPyrLK(self.prevFrame, frame, self.prevKeyPoints, None)
 
         # leave good correspondences only
         prevPoints = []
@@ -296,7 +288,7 @@ class GMC:
                 H[0, 2] *= self.downscale
                 H[1, 2] *= self.downscale
         else:
-            print("Warning: not enough matching points")
+            print('Warning: not enough matching points')
 
         # Store to next iteration
         self.prevFrame = frame.copy()
@@ -313,7 +305,7 @@ class GMC:
     def applyFile(self, raw_frame, detections=None):
         line = self.gmcFile.readline()
         tokens = line.split("\t")
-        H = np.eye(2, 3, dtype=np.float_)
+        H = np.eye(2, 3, dtype=np.float32)
         H[0, 0] = float(tokens[1])
         H[0, 1] = float(tokens[2])
         H[0, 2] = float(tokens[3])
@@ -321,4 +313,4 @@ class GMC:
         H[1, 1] = float(tokens[5])
         H[1, 2] = float(tokens[6])
 
-        return H
+        return H

supervisely/nn/tracker/{deep_sort/deep_sort → botsort/tracker}/kalman_filter.py

@@ -1,5 +1,7 @@
 # vim: expandtab:ts=4:sw=4
 import numpy as np
+import scipy.linalg
+
 
 """
 Table for the 0.95 quantile of the chi-square distribution with N degrees of
@@ -24,13 +26,13 @@ class KalmanFilter(object):
 
     The 8-dimensional state space
 
-        x, y, a, h, vx, vy, va, vh
+        x, y, w, h, vx, vy, vw, vh
 
-    contains the bounding box center position (x, y), aspect ratio a, height h,
+    contains the bounding box center position (x, y), width w, height h,
     and their respective velocities.
 
     Object motion follows a constant velocity model. The bounding box location
-    (x, y, a, h) is taken as direct observation of the state space (linear
+    (x, y, w, h) is taken as direct observation of the state space (linear
     observation model).
 
     """
@@ -56,8 +58,8 @@ class KalmanFilter(object):
         Parameters
         ----------
         measurement : ndarray
-            Bounding box coordinates (x, y, a, h) with center position (x, y),
-            aspect ratio a, and height h.
+            Bounding box coordinates (x, y, w, h) with center position (x, y),
+            width w, and height h.
 
         Returns
         -------
@@ -72,13 +74,13 @@ class KalmanFilter(object):
         mean = np.r_[mean_pos, mean_vel]
 
         std = [
+            2 * self._std_weight_position * measurement[2],
             2 * self._std_weight_position * measurement[3],
+            2 * self._std_weight_position * measurement[2],
             2 * self._std_weight_position * measurement[3],
-            1e-2,
-            2 * self._std_weight_position * measurement[3],
+            10 * self._std_weight_velocity * measurement[2],
             10 * self._std_weight_velocity * measurement[3],
-            10 * self._std_weight_velocity * measurement[3],
-            1e-5,
+            10 * self._std_weight_velocity * measurement[2],
             10 * self._std_weight_velocity * measurement[3]]
         covariance = np.diag(np.square(std))
         return mean, covariance
@@ -103,18 +105,18 @@ class KalmanFilter(object):
 
         """
         std_pos = [
+            self._std_weight_position * mean[2],
             self._std_weight_position * mean[3],
-            self._std_weight_position * mean[3],
-            1e-2,
+            self._std_weight_position * mean[2],
             self._std_weight_position * mean[3]]
         std_vel = [
+            self._std_weight_velocity * mean[2],
             self._std_weight_velocity * mean[3],
-            self._std_weight_velocity * mean[3],
-            1e-5,
+            self._std_weight_velocity * mean[2],
             self._std_weight_velocity * mean[3]]
         motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))
 
-        mean = np.dot(self._motion_mat, mean)
+        mean = np.dot(mean, self._motion_mat.T)
         covariance = np.linalg.multi_dot((
             self._motion_mat, covariance, self._motion_mat.T)) + motion_cov
 
@@ -138,9 +140,9 @@ class KalmanFilter(object):
 
         """
         std = [
+            self._std_weight_position * mean[2],
             self._std_weight_position * mean[3],
-            self._std_weight_position * mean[3],
-            1e-1,
+            self._std_weight_position * mean[2],
             self._std_weight_position * mean[3]]
         innovation_cov = np.diag(np.square(std))
 
@@ -149,6 +151,45 @@ class KalmanFilter(object):
             self._update_mat, covariance, self._update_mat.T))
         return mean, covariance + innovation_cov
 
+    def multi_predict(self, mean, covariance):
+        """Run Kalman filter prediction step (Vectorized version).
+        Parameters
+        ----------
+        mean : ndarray
+            The Nx8 dimensional mean matrix of the object states at the previous
+            time step.
+        covariance : ndarray
+            The Nx8x8 dimensional covariance matrics of the object states at the
+            previous time step.
+        Returns
+        -------
+        (ndarray, ndarray)
+            Returns the mean vector and covariance matrix of the predicted
+            state. Unobserved velocities are initialized to 0 mean.
+        """
+        std_pos = [
+            self._std_weight_position * mean[:, 2],
+            self._std_weight_position * mean[:, 3],
+            self._std_weight_position * mean[:, 2],
+            self._std_weight_position * mean[:, 3]]
+        std_vel = [
+            self._std_weight_velocity * mean[:, 2],
+            self._std_weight_velocity * mean[:, 3],
+            self._std_weight_velocity * mean[:, 2],
+            self._std_weight_velocity * mean[:, 3]]
+        sqr = np.square(np.r_[std_pos, std_vel]).T
+
+        motion_cov = []
+        for i in range(len(mean)):
+            motion_cov.append(np.diag(sqr[i]))
+        motion_cov = np.asarray(motion_cov)
+
+        mean = np.dot(mean, self._motion_mat.T)
+        left = np.dot(self._motion_mat, covariance).transpose((1, 0, 2))
+        covariance = np.dot(left, self._motion_mat.T) + motion_cov
+
+        return mean, covariance
+
     def update(self, mean, covariance, measurement):
         """Run Kalman filter correction step.
 
@@ -159,8 +200,8 @@ class KalmanFilter(object):
         covariance : ndarray
             The state's covariance matrix (8x8 dimensional).
         measurement : ndarray
-            The 4 dimensional measurement vector (x, y, a, h), where (x, y)
-            is the center position, a the aspect ratio, and h the height of the
+            The 4 dimensional measurement vector (x, y, w, h), where (x, y)
+            is the center position, w the width, and h the height of the
             bounding box.
 
         Returns
@@ -169,8 +210,6 @@ class KalmanFilter(object):
             Returns the measurement-corrected state distribution.
 
         """
-        import scipy.linalg  # pylint: disable=import-error
-
         projected_mean, projected_cov = self.project(mean, covariance)
 
         chol_factor, lower = scipy.linalg.cho_factor(
@@ -186,13 +225,11 @@ class KalmanFilter(object):
         return new_mean, new_covariance
 
     def gating_distance(self, mean, covariance, measurements,
-                        only_position=False):
+                        only_position=False, metric='maha'):
         """Compute gating distance between state distribution and measurements.
-
         A suitable distance threshold can be obtained from `chi2inv95`. If
         `only_position` is False, the chi-square distribution has 4 degrees of
         freedom, otherwise 2.
-
         Parameters
         ----------
         mean : ndarray
@@ -206,26 +243,27 @@ class KalmanFilter(object):
         only_position : Optional[bool]
             If True, distance computation is done with respect to the bounding
             box center position only.
-
         Returns
         -------
         ndarray
             Returns an array of length N, where the i-th element contains the
             squared Mahalanobis distance between (mean, covariance) and
             `measurements[i]`.
-
         """
-        import scipy.linalg  # pylint: disable=import-error
-
         mean, covariance = self.project(mean, covariance)
         if only_position:
             mean, covariance = mean[:2], covariance[:2, :2]
             measurements = measurements[:, :2]
 
-        cholesky_factor = np.linalg.cholesky(covariance)
         d = measurements - mean
-        z = scipy.linalg.solve_triangular(
-            cholesky_factor, d.T, lower=True, check_finite=False,
-            overwrite_b=True)
-        squared_maha = np.sum(z * z, axis=0)
-        return squared_maha
+        if metric == 'gaussian':
+            return np.sum(d * d, axis=1)
+        elif metric == 'maha':
+            cholesky_factor = np.linalg.cholesky(covariance)
+            z = scipy.linalg.solve_triangular(
+                cholesky_factor, d.T, lower=True, check_finite=False,
+                overwrite_b=True)
+            squared_maha = np.sum(z * z, axis=0)
+            return squared_maha
+        else:
+            raise ValueError('invalid distance metric')