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

dnt/detect/signal/detector.py

@@ -0,0 +1,317 @@
+import cv2, os
+import pandas as pd
+from pathlib import Path
+import numpy as np
+from tqdm import tqdm
+from PIL import Image
+from typing import List
+import torch
+import torch.nn as nn
+import torchvision.models as models
+import torchvision.transforms as transforms
+from torchvision.models.resnet import ResNet18_Weights
+from shared.download import download_file
+from matplotlib import pyplot as plt
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.resnet18 = models.resnet18(weights=ResNet18_Weights.DEFAULT)
+        self.resnet18.fc = nn.Linear(512, 2)
+
+    def forward(self, x):
+        return self.resnet18(x)
+    
+class SignalDetector:
+    def __init__(self, det_zones:list, model:str='ped', weights:str=None, batchsz:int=64, threshold:float=0.98, device='auto'):
+        '''
+        Detect traffic signal status
+
+        Parameters:
+        - det_zones: cropped zones for detection list[(x, y, w, h)]
+        - model: detection model, default is 'ped', 'custom'
+        - weights: path of weights, default is None
+        - batchsz: the batch size for prediction, default is 64
+        - threshold: the threshold for detection, default is 0.98
+        '''
+        
+        self.det_zones = det_zones
+
+        cwd = Path(__file__).parent.absolute()
+        if not weights:
+            if (model == 'ped'):
+                weights = os.path.join(cwd, 'weights', 'ped_signal.pt')            
+
+        if not os.path.exists(weights):
+            url = 'https://its.cutr.usf.edu/alms/downloads/ped_signal.pt'
+            download_file(url, weights)    
+        
+        self.model = Model()
+        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.model.load_state_dict(torch.load(weights))
+        self.model.to(self.device)
+
+        self.batchsz = batchsz
+        self.threshold = threshold
+
+    def detect(self, input_video:str, det_file:str=None,
+               video_index:int=None, video_tot:int=None):
+        '''
+        Parameters:
+        - input_video: the video path
+        - det_file: the file name for output
+        - video_index: the index of video
+        - video_tot: the total number of videos
+
+        Return:
+        - a dataframe with prediction results, ['frame', 'zone0', 'zone1', ...]
+        '''
+
+        # initialize the result array
+        results = np.array([])    #np.full((0, len(self.det_zones)), -1)
+        frames = np.array([])
+        zones = np.array([])
+        batch = []
+        temp_frames=[]
+
+        # open input video
+        cap = cv2.VideoCapture(input_video)
+        if not cap.isOpened():
+            print('Failed to open the video!')
+
+        tot_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        pbar = tqdm(total=tot_frames, unit=' frame')
+        if video_index and video_tot: 
+            pbar.set_description_str("Detecting signals {} of {}".format(video_index, video_tot))
+        else:
+            pbar.set_description_str("Detecting signals ")
+        
+        while cap.isOpened():
+            pos_frame = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
+            ret, frame = cap.read()
+            if not ret:
+                break
+        
+            crop_img = self.crop_zone(frame)
+            batch.append(crop_img)
+            temp_frames.append(pos_frame)
+
+            if ((pos_frame+1) % self.batchsz == 0) or (pos_frame >= tot_frames-1):
+                #batch_pred = self.predict(batch).reshape(-1, len(self.det_zones))
+                #results = np.append(results, batch_pred, axis=0)
+
+                batch_pred = self.predict(batch).flatten()
+                results = np.append(results, batch_pred, axis=0)
+                zones = np.append(zones, np.tile(np.array(list(range(len(self.det_zones)))), self.batchsz), axis=0)
+                frames = np.append(frames, np.repeat(np.array(temp_frames), len(self.det_zones)), axis=0)
+
+                batch=[]
+                temp_frames = []
+
+            pbar.update()
+
+        pbar.close()
+        cap.release()
+
+        df = pd.DataFrame(list(zip(frames, zones, results)), columns=['frame', 'signal', 'detection'])
+
+        if det_file:
+            df.to_csv(det_file, index=False)
+
+        return df
+
+    def gen_ped_interval(self, dets:pd.DataFrame, input_video:str, walk_interval:int, countdown_interval:int, 
+                         out_file:str, factor:float=0.75, video_index:int=None, video_tot:int=None):
+        '''
+        Parameters:
+        - dets: the dataframe for signal detections
+        - input_video: the video path
+        - walk_interval: the pedestrian walking interval 4s to 7s
+        - countdown_interval: the pedestrian countdown interval, cross-length(ft)/4(ft/s)
+        - factor: the factor to identify a walking signal (0, 1), 
+                    default is 0.75 (75% of frames in a sliding window is walking, then walking is identified)
+
+        Return:
+        - a dataframe of ped sigal intervals ['signal', 'walk_beg', 'walk_end', 'countdown_beg', 'countdown_end']  
+        '''
+
+         # open input video
+        cap = cv2.VideoCapture(input_video)
+        if not cap.isOpened():
+            print('Failed to open the video!')
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        cap.release()
+        
+        results = []
+        for i in range(len(self.det_zones)):
+            dets_i = dets[dets['signal']==i]
+            results.append(self.scan_walk_interval(dets_i, int(fps*walk_interval), factor, i, fps))
+
+        df = pd.concat(results, axis=0)
+        if out_file:
+            df.to_csv(out_file, index=False)
+
+        return df
+
+    def scan_walk_interval(self, dets:pd.DataFrame, window:int, factor:float, zone:int, fps:int=30,
+                           video_index:int=None, video_tot:int=None) -> pd.DataFrame:
+
+        sequence = dets['detection'].to_numpy()
+
+        if len(sequence) == 0:
+            return []    
+        
+        frame_intervals = []
+        pre_walk = False
+        tmp_cnt = 0
+        
+        pbar = tqdm(total=len(sequence)-window, unit=' frame')
+        if video_index and video_tot: 
+            pbar.set_description_str("Scanning intervals for signal {}, {} of {}".format(zone, video_index, video_tot))
+        else:
+            pbar.set_description_str("Scanning intervals for signal {}".format(zone))
+
+        for i in range(len(sequence) - window):
+            count = sum(sequence[i:i+window])
+
+            # check if the current frame can be a start of green light
+            if count >= factor * window:
+                is_walking = True
+            else:
+                is_walking = False
+
+            # if the current is green
+            # 1) if prev status is green, update the latest interval
+            # 2) if prev status is not green, append a new interval
+            if is_walking:
+                if not pre_walk:
+                    frame_intervals.append([i, i + window])
+                    tmp_cnt = 0
+                else:
+                    if count > tmp_cnt:
+                        tmp_cnt = count
+                        frame_intervals[-1] = [i, i + window]
+            
+            pre_walk = is_walking
+
+            pbar.update()
+        pbar.close()
+
+        results = []
+        for start, end in frame_intervals:
+            results.append([zone, 1, int(dets['frame'].iloc[start]), int(dets['frame'].iloc[end])]) 
+            results.append([zone, 2, int(dets['frame'].iloc[end])+1, int(dets['frame'].iloc[end]+int(10*fps))])
+    
+        df = pd.DataFrame(results, columns=['signal', 'status', 'beg_frame', 'end_frame'])
+        return df
+
+    def crop_zone(self, frame:np.ndarray)->list[np.ndarray]:
+        '''
+        Parameters:
+        - frame: frame of the video
+
+        Return:
+        - list of cropped zones
+        '''
+        crop_regions = []
+        for i, region in enumerate(self.det_zones):
+            x, y, w, h = region
+            cropped = frame[y:y + h, x:x + w]
+            cropped = Image.fromarray(cropped)
+            crop_regions.append(cropped)
+        return crop_regions
+    
+    def predict(self, batch:list):
+        """
+        Parameters:
+            batch: the cropped batch of images
+        Returns:
+            pred_array: the corresponding traffic signal predictions with format (zone0_pred, zone1_pred, ...)
+        """
+        self.model.eval()
+        sf = nn.Softmax(dim=1)
+        batchsz = len(batch)
+        num_group = len(batch[0])
+
+        # define the image transform
+        transform = transforms.Compose([
+            transforms.Resize((64, 64)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+
+        # Transform all images and stack to different groups(corresponding light)
+        transformed_crops = [[transform(image) for image in crop] for crop in batch]
+        transformed_batch = [torch.stack([transformed_crops[j][i] for j in range(batchsz)]) for i in range(num_group)]
+
+        # send to model and make predictions zone by zone
+        batch_pred = []
+        for group in transformed_batch:
+            # compute the probability of the detections
+            outputs = self.model(group.to(self.device))
+            sf_output = sf(outputs)
+
+            # predict green only if the probability > threshold
+            y_pred = (sf_output[:, 1] > self.threshold).int()
+            batch_pred.append(y_pred.data.cpu().numpy())
+
+        pred_array = np.array(batch_pred).T
+        # prob_array = np.array(batch_prob).T
+
+        return pred_array
+
+    def generate_labels(self, signals:pd.DataFrame, input_video:str, label_file:str,
+                        size_factor:float=1.5, thick:int=1, video_index:int=None, video_tot:int=None)->pd.DataFrame:
+
+         # open input video
+        cap = cv2.VideoCapture(input_video)
+        if not cap.isOpened():
+            print('Failed to open the video!')
+        tot_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        cap.release()
+
+        pbar = tqdm(total= tot_frames)
+        if video_index and video_tot:
+            pbar.set_description_str("Generating signal labeles for {} of {}".format(video_index, video_tot))
+        else:
+            pbar.set_description_str("Generating signal labels")
+        
+        results = []
+        for i in range(tot_frames):
+            for j in range(len(self.det_zones)):
+                status = 0      # no walking
+                selected = signals[(signals['beg_frame']<=i) & (signals['end_frame']>=i) & (signals['signal']==j)]
+                if len(selected) > 0:
+                    status = selected['status'].iloc[0]
+
+                x, y, w, h = self.det_zones[j]
+                cx = int(x + w/2)
+                cy = int(y + h/2)
+                r = int(max(w, h) * size_factor)
+    
+                if status == 0:
+                    results.append([i, 'circle', [(cx, cy)], (0, 0, 255), r, thick, ''])
+                elif status == 1:
+                    results.append([i, 'circle', [(cx, cy)], (0, 255, 0), r, thick, ''])
+                elif status == 2:
+                    results.append([i, 'circle', [(cx, cy)], (0, 255, 255), r, thick, ''])
+
+                pbar.update()
+
+        df = pd.DataFrame(results, columns=['frame','type','coords','color','size','thick','desc'])
+        df.sort_values(by='frame')
+        
+        if label_file:
+            df.to_csv(label_file, index=False)
+
+        return df
+

dnt/detect/yolov8/detector.py

@@ -1,27 +1,68 @@
-from ultralytics import YOLO
+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
+        model: 'yolo' (default), 'rtdetr'
+        weights: 'x' - extra-large (default), 'l' - large, 'm' - mdium, 's' - small, 'n' - nano, '.pt' - custom model weights
+        device: 'auto' (default), 'cuda', 'cpu', 'mps'
+        half: False (default), True
         '''
         # Load YOLOV8 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 == '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':
+            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
 
+        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_time:int=None, end_time:int=None, verbose:bool=False) -> pd.DataFrame:
@@ -66,7 +107,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]    
 
@@ -82,9 +124,10 @@ class Detector:
 
         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 +142,57 @@ 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()
+
+        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 +254,23 @@ 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')
+    result = detector.detect_frames('/mnt/d/videos/samples/traffic_short.mp4', [0, 1])
+    print(result)

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
+