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

dnt/__init__.py

@@ -2,4 +2,4 @@ import sys, os
 sys.path.append(os.path.dirname(__file__))
 sys.path.append(os.path.join(os.path.dirname(__file__), 'third_party/fast-reid'))
 
-__version__='0.3.1.3'
+__version__='0.3.1.7'

dnt/analysis/__init__.py

@@ -3,6 +3,6 @@ import sys
 
 sys.path.append(os.path.dirname(__file__))
 
-from .stop3 import StopAnalyzer
-from .count import Count
-from .interaction import YieldAnalyzer
+from stop3 import StopAnalyzer
+from count import Count
+from interaction2 import YieldAnalyzer

dnt/analysis/count.py

@@ -1,54 +1,73 @@
-from shapely.geometry import Point, Polygon, LineString, box
-import geopandas as gpd, pandas as pd
-import datetime
+# pylint: disable=too-many-arguments
+'''
+This script is used to count the number of tracks that passing a reference line for post-analysis
+last modified: 2021-09-30
+'''
+from shapely.geometry import Polygon, LineString
+import geopandas as gpd
+import pandas as pd
 from tqdm import tqdm
 import numpy as np
 
 class Count:
-    def __init__(self, track_fields: list[str] = ['frame', 'track_id', 'x', 'y', 'w', 'h', 'r1', 'r2', 'r3', 'r4']) -> None:
-        self.track_fields = track_fields
+    '''
+    Count tracks that passing a reference line for post-analysis
+    Methods:
+        __init__(): initialize the class
+        count_tracks_by_line(): 
+            count tracks that passing a reference line for post-analysis
+    '''
+    def __init__(self) -> None:
+        self.track_fields = ['frame', 'track_id', 'x', 'y', 'w', 'h', 'r1', 'r2', 'r3', 'r4']
 
-    def count_tracks_by_line_post(self, tracks:pd.DataFrame = None, track_file:str=None, 
-                    line: LineString = None,
-                    video_index:int = None, video_tot:int = None) -> pd.DataFrame:
+    def count_tracks_by_line(self,
+                             tracks:pd.DataFrame=None,
+                             track_file:str=None,
+                             line: LineString=None,
+                             video_index:int=None,
+                             video_tot:int=None,
+                             verbose:bool=True) -> pd.DataFrame:
         '''
         Count tracks that passing a reference line for post-analysis
             Inputs:
-                tracks - a DataFrame of tracks, [FRAME, TRACK_ID, TOPX, TOPY, WIDTH, LENGTH, RESERVED, RESERVED, RESERVED, RESERVED], 
-                         if None, read track_file
+                tracks - a DataFrame of tracks, if None (default), read track_file
                 track_file -  a txt file contains tracks, 
-                line - the reference line
+                line - the reference line, LineString(pointA, pointB)
                 video_index - the index of video for processing
                 video_tot - the total number of videos
+                verbose - if True, show progress bar
             Return:
                 A DataFrame of track_id, frame, direction, count
         '''
-
+        # Load tracks
         if tracks is None:
             tracks = pd.read_csv(track_file, header=None)
-
         if len(tracks.columns) != len(self.track_fields):
             raise Exception("The tracks format is incorrect!")
-
         tracks.columns = self.track_fields
         track_ids = tracks['track_id'].unique()
 
-        g = tracks.apply(lambda track: Polygon([(track['x'], track['y']), (track['x'] + track['w'], track['y']), 
-                                (track['x'] + track['w'], track['y'] + track['h']), (track['x'], track['y'] + track['h'])]), axis =1)
-        geo_tracks = gpd.GeoDataFrame(tracks, geometry=g)
-
-        a = np.array(line.coords[0])
-        b = np.array(line.coords[1])
-        
-        pbar = tqdm(total=len(track_ids), unit=' tracks')
-        if video_index and video_tot: 
-            pbar.set_description_str("Counting {} of {}".format(video_index, video_tot))
-        else:
-            pbar.set_description_str("Counting")
+        # Create a GeoDataFrame of tracks
+        geo = tracks.apply(lambda track: Polygon([(track['x'], track['y']),
+                                                (track['x'] + track['w'], track['y']),
+                                                (track['x'] + track['w'], track['y'] + track['h']),
+                                                (track['x'], track['y'] + track['h'])]),
+                                                axis =1)
+        geo_tracks = gpd.GeoDataFrame(tracks, geometry=geo)
+        point_a = np.array(line.coords[0])
+        point_b = np.array(line.coords[1])
 
+        # Interate through all tracks
+        if verbose:
+            pbar = tqdm(total=len(track_ids), unit=' tracks')
+            if video_index and video_tot:
+                pbar.set_description_str("Counting {} of {}".format(video_index, video_tot))
+            else:
+                pbar.set_description_str("Counting")
         intersected_tracks = []
         intersected_frames = []
-        intersected_direct = []   
+        intersected_direct = []
+
         for track_id in track_ids:
             selected = geo_tracks.loc[(geo_tracks['track_id']==track_id)].copy()
             if len(selected)>0:
@@ -60,19 +79,19 @@ class Count:
                         intersected_tracks.append(track_id)
                         frame_pos = int(len(intersected)/2)
                         intersected_frames.append(intersected['frame'].values[frame_pos])
-                        
+
                         # center point of the first frame
-                        p = np.array((intersected.iloc[0]['x'] + intersected.iloc[0]['w']/2,
+                        point_c = np.array((intersected.iloc[0]['x'] + intersected.iloc[0]['w']/2,
                                     intersected.iloc[0]['y'] + intersected.iloc[0]['h']/2))
                         d = 2  # right
-                        if np.cross(p-a, b-a) < 0:
+                        if np.cross(point_c-point_a, point_b-point_a) < 0:
                             d = 1 # left
-
                         intersected_direct.append(d)
-        
-            pbar.update()
-
-        pbar.close()
+            if verbose:
+                pbar.update()
+                
+        if verbose:
+            pbar.close()
 
         results = pd.DataFrame(
             {
@@ -86,5 +105,3 @@ class Count:
         results['count'] = results.index + 1
 
         return results
-    
-    

dnt/analysis/{interaction.py → interaction2.py}

@@ -1,21 +1,36 @@
+"""
+This module is used to analyze the interaction between two tracks.
+    Author: Zhenyu Wang
+    Date: 2023/10/20
+    Email: wonstran@hotmail.com
+    Updates:
+        2025/03/19: Add the function of nearmiss based on PET
+        2023/10/20: Initial version
+"""
 from shapely.geometry import Point, Polygon, LineString, box
 from shapely import intersection, distance, intersects
 from shapelysmooth import taubin_smooth, chaikin_smooth, catmull_rom_smooth
 import geopandas as gpd, pandas as pd
 from tqdm import tqdm
 import numpy as np
-from matplotlib import pyplot as plt
 from dnt.label.labeler2 import Labeler
 import os
-from ..track import Tracker
 
 class YieldAnalyzer:
-    def __init__(self, waiting_dist_p:int=600, waiting_dist_y:int=300, 
-                 leading_p:bool=False, leading_axis_p:str='y',
-                 leading_y:bool=True, leading_axis_y:str='x',
-                 yield_gap:int=10, fps:int=30, ref_point='bc', 
-                 ref_offset:tuple=(0,0), filter_buffer:int=0,
-                 p_zone:Polygon=None, y_zone:Polygon=None) -> None:
+    def __init__(self,
+                 waiting_dist_p:int=600,
+                 waiting_dist_y:int=300, 
+                 leading_p:bool=False,
+                 leading_axis_p:str='y',
+                 leading_y:bool=True,
+                 leading_axis_y:str='x',
+                 yield_gap:int=10,
+                 fps:int=30,
+                 ref_point='bc', 
+                 ref_offset:tuple=(0,0),
+                 filter_buffer:int=0,
+                 p_zone:Polygon=None,
+                 y_zone:Polygon=None) -> None:
         '''
         Parameters:
             threshold: the hyperparameter to determine if a yield event (frame difference <=yield_gap*fps), default is 3 seconds

dnt/analysis/stop3.py

@@ -6,11 +6,15 @@ from tqdm import tqdm
 from cython_bbox import bbox_overlaps
 import numpy as np
 #import mapply
-from ..filter import Filter
+import sys
+import os
+sys.path.append(os.path.dirname(__file__))
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+from filter import Filter
 from matplotlib import pyplot as plt
 import random, os
-from ..label.labeler2 import Labeler
-from ..engine.bbox_iou import ious
+from label.labeler2 import Labeler
+from engine.bbox_iou import ious
 
 class StopAnalyzer():
     def __init__(self, 

dnt/detect/signal/detector.py

@@ -12,19 +12,27 @@ import torchvision.transforms as transforms
 from torchvision.models.resnet import ResNet18_Weights
 from shared.download import download_file
 from matplotlib import pyplot as plt
+from ultralytics import YOLO
 
 class Model(nn.Module):
-    def __init__(self):
+    def __init__(self, num_class=2):
         super(Model, self).__init__()
 
         self.resnet18 = models.resnet18(weights=ResNet18_Weights.DEFAULT)
-        self.resnet18.fc = nn.Linear(512, 2)
+        self.resnet18.fc = nn.Linear(512, num_class)
 
     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'):
+    def __init__(self,
+                 det_zones:list,
+                 model:str='ped',
+                 weights:str=None,
+                 batchsz:int=64,
+                 num_class:int=2,
+                 threshold:float=0.98,
+                 device='auto'):
         '''
         Detect traffic signal status
 
@@ -33,6 +41,7 @@ class SignalDetector:
         - model: detection model, default is 'ped', 'custom'
         - weights: path of weights, default is None
         - batchsz: the batch size for prediction, default is 64
+        - num_class: the number of classes, default is 2
         - threshold: the threshold for detection, default is 0.98
         '''
         
@@ -47,7 +56,7 @@ class SignalDetector:
             url = 'https://its.cutr.usf.edu/alms/downloads/ped_signal.pt'
             download_file(url, weights)    
         
-        self.model = Model()
+        self.model = Model(num_class)
         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():

dnt/detect/timestamp.py

@@ -0,0 +1,105 @@
+import easyocr
+import cv2
+import numpy as np
+import dateparser
+from datetime import datetime
+
+class TimestampExtractor:
+    def __init__(self,
+                 zone:np.ndarray=None,
+                 allowlist='0123456789-:/'):
+        """
+        args:
+            zone (np.ndarray): The coordinates of the timestamp area in the image.
+            allowlist (str): Characters to allow in the OCR process.
+        """
+        self.zone = zone
+        self.reader = easyocr.Reader(['en'])
+        self.allowlist = allowlist
+
+    def extract_timestamp(self,
+                          img:np.ndarray,
+                          gray:bool=False,
+                          ) -> datetime:
+        """
+        Extract timestamp from the image using OCR.
+        args:
+            img (np.ndarray): The input image from which to extract the timestamp.
+        returns:
+            datetime: The extracted timestamp.
+        """
+        # Ensure the zone is a numpy array
+        if self.zone is None:
+            crop_img = img
+        else:
+            x1, y1, x2, y2 = self.zone[0][0], self.zone[0][1], self.zone[2][0], self.zone[2][1]
+            crop_img = img[y1:y2, x1:x2]
+
+        if gray:
+            crop_img = cv2.cvtColor(crop_img, cv2.COLOR_BGR2GRAY)
+        else:
+            crop_img = crop_img
+
+        result = self.reader.readtext(crop_img, detail=0, allowlist='0123456789-:/')
+        dt = dateparser.parse(" ".join(result))
+        return dt
+    
+    def extract_timestamp_video(self,
+                                video_path:str,
+                                frame:int=0,
+                                gray:bool=False,
+                                ) -> datetime:
+        """
+        Extract timestamp from a video file using OCR.
+        args:
+            video_path (str): The path to the video file.
+            frame (int): The frame index to extract the timestamp from.
+            gray (bool): Whether to convert the image to grayscale before OCR.
+        returns:
+            datetime: The extracted timestamp.
+        """
+        cap = cv2.VideoCapture(video_path)
+        cap.set(cv2.CAP_PROP_POS_FRAMES, frame)
+        ret, img = cap.read()
+        cap.release()
+
+        if not ret:
+            raise ValueError("Could not read video file")
+
+        return self.extract_timestamp(img, gray=gray)
+    
+    def extract_timestamp_video_auto(self,
+                                    video_path:str,
+                                    gray:bool=False,
+                                    ) -> tuple:
+        """
+        Extract the initial timestamp from the first readable frame in a video file using OCR.
+        args:
+            video_path (str): The path to the video file.
+            gray (bool): Whether to convert the image to grayscale before OCR.
+        returns:
+            dt, frame: The extracted timestamp and the frame index.
+        """
+        cap = cv2.VideoCapture(video_path)
+        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
+        tot = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+        for frame in range(0, tot-1):
+            ret, img = cap.read()
+            if not ret:
+                break
+            dt = self.extract_timestamp(img, gray=gray)
+            if dt:
+                cap.release()
+                return dt, frame
+        
+        cap.release()
+        return None, None    
+
+if __name__ == "__main__":
+    img = cv2.imread('/mnt/d/videos/sample/frames/007_day.png')
+    zone = np.array([[62, 1003], [543, 1000], [547, 1059], [67, 1061]])
+    extractor = TimestampExtractor(zone=zone)
+    timestamp = extractor.extract_timestamp(img)
+    print("Extracted Timestamp:")
+    print(timestamp)

dnt/detect/yolov8/detector.py

@@ -1,3 +1,6 @@
+'''
+This module is a wrapper for ultralytics YOLOV8 and RTDETR models.
+'''
 from ultralytics import YOLO, RTDETR
 import cv2
 import pandas as pd
@@ -7,21 +10,38 @@ from pathlib import Path
 import torch
 
 class Detector:
-    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):
+    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):
         '''
-        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
+        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()
         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')
@@ -39,7 +59,7 @@ class Detector:
         else:
             raise Exception('Invalid detection model type!')
         
-        if model == 'yolo':
+        if model == 'yolo' or model == 'yolo11':
             self.model = YOLO(model_path)
         elif model == 'rtdetr':
             self.model = RTDETR(model_path)
@@ -63,10 +83,27 @@ class Detector:
     
         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:
-     
+    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():
@@ -76,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
@@ -93,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():
@@ -120,9 +156,11 @@ 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()
 
@@ -179,7 +217,11 @@ class Detector:
             pbar.close()
         cap.release()
 
-        df = pd.concat(results)
+        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)
@@ -271,6 +313,7 @@ class Detector:
 
 if __name__=='__main__':
 
-    detector = Detector()
-    result = detector.detect_frames('/mnt/d/videos/samples/traffic_short.mp4', [0, 1])
+    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}")