dnt 0.2.4__py3-none-any.whl → 0.3.1.7__py3-none-any.whl

dnt/detect/yolov8/detector.py

@@ -1,31 +1,109 @@
-from ultralytics import YOLO
+'''
+This module is a wrapper for ultralytics YOLOV8 and RTDETR models.
+'''
+from ultralytics import YOLO, RTDETR
 import cv2
 import pandas as pd
 from tqdm import tqdm
 import random, os
 from pathlib import Path
+import torch
 
 class Detector:
-    def __init__(self, yolo:str='x', conf:float=0.25, nms:float=0.7, max_det:int=300):
+    def __init__(self,
+                 model:str='yolo',
+                 weights:str='x',
+                 conf:float=0.25,
+                 nms:float=0.7,
+                 max_det:int=300, 
+                 device:str='auto',
+                 half:bool=False):
         '''
-        yolo:  x - yolov8x.pt 
-                l - yolov8l.pt
-                m - yolov8m.pt
-                s - yolov8s.pt
-                n - yolov8n.pt
+        Initialize the Detector class.
+        Parameters:
+            model: 'yolo' (default), 'rtdetr'
+            weights: 'x' - extra-large (default), 'l' - large, 'm' - mdium, 's' - small, 'n' - nano, '.pt' - custom model weights
+            conf: 0.25 (default) - confidence threshold
+            nms: 0.7 (default) - non-max suppression threshold
+            max_det: 300 (default) - maximum number of detections per image
+            device: 'auto' (default), 'cuda', 'cpu', 'mps'
+            half: False (default), True
         '''
-        # Load YOLOV8 model
+        # Load model
         cwd = Path(__file__).parent.absolute()
-        model_path = os.path.join(cwd, 'models/yolov8'+yolo+'.pt')
-        self.model = YOLO(model_path)     
+        if model == 'yolo':
+            if weights in ['x', 'l', 'm', 's', 'n']:
+                model_path = os.path.join(cwd, 'models/yolov8'+weights+'.pt')
+            elif ".pt" in weights:
+                model_path = os.path.join(cwd, 'models/'+weights)
+        elif model == 'yolo11':
+            if weights in ['x', 'l', 'm', 's', 'n']:
+                model_path = os.path.join(cwd, 'models/yolo11'+weights+'.pt')
+            elif ".pt" in weights:
+                model_path = os.path.join(cwd, 'models/'+weights)
+
+        elif model == 'rtdetr':
+            if weights in ['x', 'l']:
+                model_path = os.path.join(cwd, 'models/rtdetr-'+weights+'.pt')
+            elif ".pt" in weights:
+                model_path = os.path.join(cwd, 'models/'+weights)
+            else:
+                model_path = os.path.join(cwd, 'models/rtdetr-x.pt')
+        #elif model == 'nas':
+        #    if weights in ['l', 'm', 's']:
+        #        model_path = os.path.join(cwd, 'models/yolo_nas_'+weights+'.pt')
+        #    elif ".pt" in weights:
+        #        model_path = os.path.join(cwd, 'models/'+weights)
+        #    else:
+        #        model_path = os.path.join(cwd, 'models/yolo_nas_l.pt')
+        else:
+            raise Exception('Invalid detection model type!')
+        
+        if model == 'yolo' or model == 'yolo11':
+            self.model = YOLO(model_path)
+        elif model == 'rtdetr':
+            self.model = RTDETR(model_path)
+        #elif model == 'nas':
+        #    self.model = NAS(model_path)
+        else:
+            raise Exception('Invalid model weights!')  
+        
         self.conf = conf
         self.nms = nms
         self.max_det = max_det
 
-    def detect(self, input_video:str, iou_file:str=None,
-               video_index:int=None, video_tot:int=None,
-               start_time:int=None, end_time:int=None, verbose:bool=False) -> pd.DataFrame:
-     
+        if device == 'auto':
+            self.device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
+        elif device == 'cuda' and torch.cuda.is_available():
+            self.device = 'cuda'
+        elif device == 'mps' and torch.backends.mps.is_available():
+            self.device = 'mps'
+        else:
+            self.device = 'cpu'        
+    
+        self.half = half
+
+    def detect(self,
+               input_video:str,
+               iou_file:str=None,
+               video_index:int=None,
+               video_tot:int=None,
+               start_frame:int=None,
+               end_frame:int=None,
+               verbose:bool=True) -> pd.DataFrame:
+        '''
+        Detect objects in a video file.
+        Parameters:
+            input_video: path to the input video file
+            iou_file: path to the output file
+            video_index: index of the video in the batch
+            video_tot: total number of videos in the batch
+            start_frame: start frame number
+            end_frame: end frame number
+            verbose: True (default), False
+        Returns:
+            df: pandas DataFrame containing the detection results
+        '''
         # open input video
         cap = cv2.VideoCapture(input_video)
         if not cap.isOpened():
@@ -35,16 +113,14 @@ class Detector:
         results = []
 
         video_fps = int(cap.get(cv2.CAP_PROP_FPS))                    #original fps
-        if start_time:
-            start_frame = int(video_fps * start_time)
-            if start_frame > int(cap.get(cv2.CAP_PROP_FRAME_COUNT))-1:
+        if start_frame:
+            if (start_frame > int(cap.get(cv2.CAP_PROP_FRAME_COUNT))-1) or (start_frame < 0):
                 start_frame = 0    
         else:
             start_frame = 0
 
-        if end_time:
-            end_frame = int(video_fps * end_time)
-            if end_frame > int(cap.get(cv2.CAP_PROP_FRAME_COUNT))-1:
+        if end_frame:
+            if (end_frame > int(cap.get(cv2.CAP_PROP_FRAME_COUNT))-1) or (end_frame < 0):
                 end_frame = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))-1    
         else:
             end_frame = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))-1
@@ -52,12 +128,13 @@ class Detector:
         if start_frame>=end_frame:
             raise Exception('The given start time exceeds the end time!')
         
-        frame_total = end_frame - start_frame      
-        pbar = tqdm(total=frame_total, unit=" frames")
-        if video_index and video_tot: 
-            pbar.set_description_str("Detecting {} of {}".format(video_index, video_tot))
-        else:
-            pbar.set_description_str("Detecting ")
+        frame_total = end_frame - start_frame
+        if verbose:      
+            pbar = tqdm(total=frame_total, unit=" frames")
+            if video_index and video_tot: 
+                pbar.set_description_str("Detecting {} of {}".format(video_index, video_tot))
+            else:
+                pbar.set_description_str("Detecting ")
 
         cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
         while cap.isOpened():
@@ -66,7 +143,8 @@ class Detector:
             if (not ret) or (pos_frame>end_frame):
                 break
             
-            detects = self.model.predict(frame, verbose=False, conf=self.conf, iou=self.nms, max_det=self.max_det)   
+            detects = self.model.predict(frame, verbose=False, conf=self.conf, iou=self.nms, max_det=self.max_det, 
+                                         device=self.device, half=self.half)   
             if len(detects)>0:
                 detect = detects[0]    
 
@@ -78,13 +156,16 @@ class Detector:
                 result.insert(1, 'res', -1)
                 results.append(result)
 
-            pbar.update()
+            if verbose:
+                pbar.update()
 
-        pbar.close()
+        if verbose:
+            pbar.close()
         cap.release()
-        cv2.destroyAllWindows()
+        #cv2.destroyAllWindows()
 
-        df = pd.concat(results)
+        df = pd.concat(d for d in results if not d.empty) # remove empty dataframes
+        #df = pd.concat(results)
         df['x'] = round(df['x'], 1)
         df['y'] = round(df['y'], 1)
         df['w'] = round(df['x2']-df['x'], 0)
@@ -99,9 +180,61 @@ class Detector:
         
         return df
 
+    def detect_frames(self, input_video:str, frames:list[int], verbose:bool=False)->pd.DataFrame:
+        # open input video
+        cap = cv2.VideoCapture(input_video)
+        if not cap.isOpened():
+            print('Failed to open the video!')
+
+        random.seed(3)  # deterministic bbox colors
+        results = []
+
+        if verbose:
+            pbar = tqdm(total=len(frames), unit=" frames")        
+        for pos_frame in frames:
+            cap.set(cv2.CAP_PROP_POS_FRAMES, pos_frame)
+            ret, frame = cap.read()
+            if not ret:
+                break
+            
+            detects = self.model.predict(frame, verbose=False, conf=self.conf, iou=self.nms, max_det=self.max_det, 
+                                         device=self.device, half=self.half)   
+            if len(detects)>0:
+                detect = detects[0]    
+
+                xyxy = pd.DataFrame(detect.boxes.xyxy.tolist(), columns=['x', 'y', 'x2', 'y2'])
+                cls = pd.DataFrame(detect.boxes.cls.tolist(), columns=['class'])
+                conf = pd.DataFrame(detect.boxes.conf.tolist(), columns = ['conf'])
+                result = pd.concat([xyxy, conf, cls], axis=1)
+                result.insert(0, 'frame', pos_frame)
+                result.insert(1, 'res', -1)
+                results.append(result)
+
+            if verbose:
+                pbar.update()
+
+        if verbose:
+            pbar.close()
+        cap.release()
+
+        if len(results) == 0:
+            return pd.DataFrame(columns=['frame', 'res', 'x', 'y', 'w', 'h', 'conf', 'class'])
+        
+        df = pd.concat(d for d in results if not d.empty) # remove empty dataframes
+        #df = pd.concat(results)
+        df['x'] = round(df['x'], 1)
+        df['y'] = round(df['y'], 1)
+        df['w'] = round(df['x2']-df['x'], 0)
+        df['h'] = round(df['y2']-df['y'], 0)
+        df['class'] = round(df['class'], 0)
+        df['conf'] = round(df['conf'], 2)
+        df = df[['frame', 'res', 'x', 'y', 'w', 'h', 'conf', 'class']].reset_index(drop=True)
+
+        return df
+
     def detect_v8(self, input_video:str, iou_file:str=None, save:bool=False, verbose:bool=False, show:bool=False):
         detects = self.model.predict(input_video, 
-                        verbose=verbose, stream=True, save=save, show=show)
+                        verbose=verbose, stream=True, save=save, show=show, device=self.device, half=self.half)
         
         if iou_file:
             results = []
@@ -163,14 +296,24 @@ class Detector:
         return results
 
     @staticmethod
-    def get_fps(input_video:str)->float:
-        cap = cv2.VideoCapture(input_video)
+    def get_fps(video:str)->float:
+        cap = cv2.VideoCapture(video)
         if not cap.isOpened():
             print('Failed to open the video!')
         fps = cap.get(cv2.CAP_PROP_FPS)
         return fps
+    
+    @staticmethod
+    def get_frames(video:str)->int:
+        cap = cv2.VideoCapture(video)
+        if not cap.isOpened():
+            print('Failed to open the video!')
+        frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        return frames
 
 if __name__=='__main__':
 
-    detector = Detector()
-    detector.detect('/mnt/f/videos/samples/traffic_short.mp4', '/mnt/f/videos/samples/traffic_short_iou.txt')
+    detector = Detector(model='yolo11')
+    result = detector.detect('/mnt/d/videos/sample/traffic.mp4', verbose=True)
+
+    print(result)

dnt/detect/yolov8/segmentor.py

@@ -4,15 +4,18 @@ import pandas as pd
 from tqdm import tqdm
 import random, os, json
 from pathlib import Path
+import torch
+from shapely.geometry import Polygon as Pol
 
 class Segmentor:
-    def __init__(self, model:str='yolov8x-seg.pt', conf:float=0.25, nms:float=0.7, max_det:int=300, half:bool=False):
+    def __init__(self, model:str='yolov8m-seg.pt', conf:float=0.25, nms:float=0.7, max_det:int=300, device:str='auto', half:bool=False):
         '''
         yolo:  x - yolov8x-seg.pt 
                 l - yolov8l-seg.pt
                 m - yolov8m-seg.pt
                 s - yolov8s-seg.pt
                 n - yolov8n-seg.pt
+        
         '''
         # Load YOLOV8 model
         cwd = Path(__file__).parent.absolute()
@@ -21,6 +24,17 @@ class Segmentor:
         self.conf = conf
         self.nms = nms
         self.max_det = max_det
+
+        if device == 'auto':
+            self.device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
+        elif device == 'cuda' and torch.cuda.is_available():
+            self.device = 'cuda'
+        elif device == 'mps' and torch.backends.mps.is_available():
+            self.device = 'mps'
+        else:
+            self.device = 'cpu'        
+    
+        self.half = half
         self.half = half
 
     def segment(self, input_video:str, mask_file:str=None,
@@ -106,8 +120,52 @@ class Segmentor:
                video_index=None, video_tot=None,
                start_frame=frame, end_frame=frame, 
                verbose=False)
+    
+    def segment_zone(self, input_video:str, frame:int=None, zone:list=None) -> list:
+        '''
+        Segment the video in a specific zone
+        input_video: Path to video
+        frame: Frame number to segment
+        zone: Polygon zone to segment
+        Return:
+            List of segmented objects in the zone
+        '''        
+    
+        # open input video
+        cap = cv2.VideoCapture(input_video)
+        if not cap.isOpened():
+                print('Failed to open the video!')
 
+        if frame <0 or frame > int(cap.get(cv2.CAP_PROP_FRAME_COUNT))-1:
+            raise Exception('The given frame is out of range!')       
+
+        cap.set(cv2.CAP_PROP_POS_FRAMES, frame)
+        ret, frame = cap.read()
+        if ret:
+            x1, y1, x2, y2 = zone
+            cropped_frame = frame[y1:y2, x1:x2]           
+            detects = self.model.predict(cropped_frame, verbose=False, conf=self.conf, iou=self.nms, max_det=self.max_det, half=self.half)   
+            if len(detects)>0:
+                if detects[0].masks is not None:
+                    xy = list(map(lambda x: x.tolist(), detects[0].masks.xy))
+                    print(xy)
+                    input('.')
+        cap.release()
+                
+        
 if __name__=="__main__":
     
     segmentor = Segmentor()
-    segmentor.segment('/mnt/d/videos/samples/traffic_short.mp4', '/mnt/d/videos/samples/traffic_short_mask.json', start_frame=1, end_frame=1)
+    #segmentor.segment('/mnt/d/videos/samples/traffic_short.mp4', '/mnt/d/videos/samples/traffic_short_mask.json', start_frame=1, end_frame=1)
+    input_video = '/mnt/d/videos/hyde/after_2024-08-15_nw.mp4'
+    tracks = pd.read_csv('/mnt/d/videos/hyde/tracks/after_2024-08-15_nw_track_veh_rt.txt')
+    tracks.columns = ['frame', 'track', 'x', 'y', 'w', 'h', 'r1', 'r2', 'r3', 'r4']
+    dets = tracks[tracks['track'] == 678]
+    dets.sort_values(by='frame', inplace=True)
+    
+    for index, det in dets.iterrows():
+        zone = (int(det['x']), int(det['y']), int(det['x'] + det['w']), int(det['y'] + det['h']))
+        frame = int(det['frame'])
+
+        segmentor.segment_zone(input_video, frame, zone)  # Segment the specific zone
+        print(f"Segmented frame {frame}")

dnt/engine/__init__.py

@@ -0,0 +1,8 @@
+import os, sys
+
+sys.path.append(os.path.dirname(__file__))
+
+from .bbox_interp import interpolate_bbox, interpolate_bboxes
+from .bbox_iou import ious
+from .cluster import cluster_by_gap
+from .iob import iob, iobs

dnt/engine/bbox_interp.py

@@ -0,0 +1,83 @@
+# bbox_interpolation.pyx
+
+import numpy as np
+from scipy.interpolate import CubicSpline, interp1d
+
+# Define the function signature
+def interpolate_bbox(boxes:np.ndarray, frames:np.ndarray, target_frame:int, method='cubic')->np.ndarray:
+    """
+    Interpolates bounding boxes using cubic splines.
+    
+    Parameters:
+    - boxes: a (n, 4) array of bounding box coordinates (x, y, width, height).
+    - frames: a (n,) array of frame indices corresponding to the bounding boxes.
+    - target_frame: the frame index at which to interpolate the bounding box.
+    - method: the spline function, default is 'cubic', 'nearest', 'linear' 
+    Returns:
+    - A 1d array (x, y, width, height) representing the interpolated bounding box.
+    """
+    n = boxes.shape[0]
+    
+    if n != frames.shape[0]:
+        raise ValueError("Length of boxes and frames must be equal.")
+    
+    if n == 0:
+        raise ValueError("Input arrays must not be empty.")
+    
+    if target_frame < frames[0] or target_frame > frames[-1]:
+        raise ValueError("Target frame is out of bounds.")
+    
+    # Unpack the boxes into separate arrays for x, y, width, and height
+    x = boxes[:, 0]
+    y = boxes[:, 1]
+    w = boxes[:, 2]
+    h = boxes[:, 3]
+    
+    # Create the cubic splines for each parameter
+    if method == 'cubic':
+        spline_x = CubicSpline(frames, x)
+        spline_y = CubicSpline(frames, y)
+        spline_w = CubicSpline(frames, w)
+        spline_h = CubicSpline(frames, h)
+    elif method == 'nearest':
+        spline_x = interp1d(frames, x, kind='nearest')
+        spline_y = interp1d(frames, y, kind='nearest')
+        spline_w = interp1d(frames, w, kind='nearest')
+        spline_h = interp1d(frames, h, kind='nearest')
+    else:
+        spline_x = interp1d(frames, x, kind='linear')
+        spline_y = interp1d(frames, y, kind='linear')
+        spline_w = interp1d(frames, w, kind='linear')
+        spline_h = interp1d(frames, h, kind='linear')
+        
+    # Evaluate the splines at the target frame
+    x_t = int(spline_x(target_frame))
+    y_t = int(spline_y(target_frame))
+    w_t = int(spline_w(target_frame))
+    h_t = int(spline_h(target_frame))
+    
+    
+    return np.array([x_t, y_t, w_t, h_t])
+
+def interpolate_bboxes(boxes:np.ndarray, frames:np.ndarray, target_frames:np.ndarray, method='cubic')->np.ndarray:
+    """
+    Interpolates bounding boxes using cubic splines.
+    
+    Parameters:
+    - boxes: a (n, 4) array of bounding box coordinates (x, y, width, height).
+    - frames: a (n,) array of frame indices corresponding to the bounding boxes.
+    - target_frames: the frame indexes at which to interpolate the bounding boxes.
+    - method: the spline function, default is 'cubic', 'nearest', 'linear' 
+
+    Returns:
+    - List of tuple (x, y, width, height) representing the interpolated bounding boxes.
+    """
+
+    n_frames = target_frames.shape[0]
+    results = []
+
+    for i in range(n_frames):
+        target_frame = target_frames[i]
+        results.append(interpolate_bbox(boxes, frames, target_frame, method))
+    
+    return np.vstack(results)

dnt/engine/bbox_iou.py

@@ -0,0 +1,20 @@
+from cython_bbox import bbox_overlaps
+import numpy as np
+
+def ious(atlbrs, btlbrs):
+    """
+    Compute cost based on IoU
+    :type atlbrs: list[tlbr] | np.ndarray
+    :type atlbrs: list[tlbr] | np.ndarray
+
+    :rtype ious np.ndarray
+    """
+    ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float64)
+    if ious.size == 0:
+        return ious
+
+    ious = bbox_overlaps(
+        np.ascontiguousarray(atlbrs, dtype=np.float64),
+        np.ascontiguousarray(btlbrs, dtype=np.float64))
+
+    return ious

dnt/engine/cluster.py

@@ -0,0 +1,31 @@
+import numpy as np
+
+def cluster_by_gap(arr:np.ndarray, gap:int)->list:
+    '''
+    Parameters:
+    - arr: a 1d array for clustering
+    - gap: the threshold to break (> gap)
+
+    Returns:
+    - list of clustered numbers (lists)
+    '''    
+    # Sort the array to process it in order
+    arr.sort()
+    clusters = []
+    current_cluster = [arr[0]]
+
+    # Iterate through the sorted list
+    for i in range(1, len(arr)):
+        if arr[i] - arr[i-1] <= gap:
+            # If the gap condition is met, add to the current cluster
+            current_cluster.append(arr[i])
+        else:
+            # If the gap condition is not met, start a new cluster
+            clusters.append(current_cluster)
+            current_cluster = [arr[i]]
+
+    # Add the last cluster
+    clusters.append(current_cluster)
+
+    return clusters
+

dnt/engine/iob.py

@@ -0,0 +1,66 @@
+import numpy as np
+
+def iob(box_a:np.ndarray, box_b:np.ndarray)->tuple:
+    """
+    Calculate the Intersection over Box (IoB) between two bounding boxes.
+
+    Parameters:
+        box_a -- numpy arrays of shape (4,) containing [l, t, w, h]
+        box_b -- numpy arrays of shape (4,) containing [l, t, w, h], target bounding box
+
+    Returns:
+    iob, iob -- float, the IoB values for box_a and box_b
+    """
+    # Determine the (x, y)-coordinates of the intersection rectangle
+    x_l = max(box_a[0], box_b[0])
+    y_t = max(box_a[1], box_b[1])
+    x_r = min(box_a[0]+box_a[2], box_b[0]+box_b[2])
+    y_b = min(box_a[1]+box_a[3], box_b[1]+box_b[3])
+    
+    # Compute the area of intersection rectangle
+    inter_width = max(0, x_r - x_l)
+    inter_height = max(0, y_b - y_t)
+    inter_area = inter_width * inter_height
+    
+    # Compute the area of both the prediction and ground-truth rectangles
+    box_a_area = box_a[2] * box_a[3]
+    box_b_area = box_b[2] * box_b[3]
+    
+    # Compute the intersection over union by taking the intersection area and dividing it by the sum of prediction + ground-truth areas - the intersection area
+    iob_a = inter_area / box_a_area if box_a_area != 0 else 0.0
+    iob_b = inter_area / box_b_area if box_b_area != 0 else 0.0
+    
+    return iob_a, iob_b
+
+def iobs(alrbs:np.ndarray, blrbs:np.ndarray)->np.ndarray:
+    """
+    Calculate the IoU matrix for multiple bounding boxes.
+
+    Parameters:
+    boxes -- numpy array of shape (N, 4), where N is the number of boxes, and each row is [x1, y1, x2, y2]
+
+    Returns:
+    iou_matrix -- numpy array of shape (N, N), where iou_matrix[i, j] is the IoU between boxes[i] and boxes[j]
+    """
+    num_a = alrbs.shape[0]
+    num_b = blrbs.shape[0]
+    iobs_a = np.zeros((num_a, num_b))
+    iobs_b = np.zeros((num_a, num_b))
+
+    for i in range(num_a):
+        for j in range(num_b):
+            iobs_a[i, j], iobs_b[i, j] = iob(alrbs[i,:], blrbs[j,:])
+            
+    return iobs_a, iobs_b
+
+if __name__ == '__main__':
+    # Example usage
+    boxes = np.array([
+        [0, 0, 1, 1],
+        [0.5, 0.5, 1.5, 1.5],
+        [1, 1, 2, 2]
+    ])
+
+    matrix_a, matrix_b = iobs(boxes, boxes)
+    print(matrix_a)
+    print(matrix_b)