dnt 0.2.1__py3-none-any.whl → 0.3.1.8__py3-none-any.whl

dnt/analysis/stop.py

@@ -1,9 +1,17 @@
 from shapely.geometry import Point, Polygon, LineString, box
 import geopandas as gpd, pandas as pd
+import datetime
 from tqdm import tqdm
 
 class StopAnalyzer():
-    def __init__(self, h_coords, v_coords, event_dict, stop_iou=0.97, frame_buffer=5, verbose=True):
+    def __init__(self, 
+                 h_coords:list[list[tuple[int, int], tuple[int, int]]], 
+                 v_coords:list[list[tuple[int, int], tuple[int, int]]], 
+                 event_dict:dict[int, list[int]], 
+                 stop_iou:float=0.97, 
+                 frame_buffer:int=5, 
+                 verbose:bool=True):
+        
         self.hzones = StopAnalyzer.gen_zones(h_coords)
         self.vzones = StopAnalyzer.gen_zones(v_coords)
         self.event_dict = event_dict
@@ -11,13 +19,17 @@ class StopAnalyzer():
         self.frame_buffer = frame_buffer
         self.verbose = verbose
     
-    def analysis(self, track_file, result_file=None, output_file=None, video_index=None, video_tot=None):
+    def analysis(self, track_file:str, 
+                 result_file:str=None, 
+                 output_file:str=None, 
+                 video_index:int=None, 
+                 video_tot:int=None):
         
         tracks = pd.read_csv(track_file, header=None, sep=',')
 
-        tracks = self.__stop_scan(tracks, video_index, video_tot)
-        tracks = self.__event_identify(tracks, video_index, video_tot)
-        results = self.__event_count(tracks, video_index, video_tot)
+        tracks = self.stop_scan(tracks, video_index, video_tot)
+        tracks = self.event_identify(tracks, video_index, video_tot)
+        results = self.event_count(tracks, video_index, video_tot)
 
         if result_file:
             results.to_csv(result_file, index=False)
@@ -25,17 +37,17 @@ class StopAnalyzer():
         if output_file:
             tracks.to_csv(output_file, header=None, index=False)
     
-    def __stop_scan(self, tracks, video_index, video_tot):
+    def stop_scan(self, tracks:pd.DataFrame, video_index:int=None, video_tot:int=None)->pd.DataFrame:
         vehicles = tracks[1].unique()
 
         pbar = tqdm(total=len(vehicles), unit=' tracks')
         if video_index and video_tot:
-            pbar.set_description_str("Scan stops {} of {}".format(video_index, video_tot))
+            pbar.set_description_str("Scanning stops {} of {}".format(video_index, video_tot))
         else:
-            pbar.set_description_str("Scan stops ")
+            pbar.set_description_str("Scanning stops ")
 
         for vehicle in vehicles:
-            track = tracks[tracks[1] == vehicle].sort_values(by=0)
+            track = tracks[tracks[1] == vehicle].sort_values(by=0) # sort by frame in ascending
 
             for i in range(self.frame_buffer, len(track)):
 
@@ -46,17 +58,22 @@ class StopAnalyzer():
                 bb1 = [track.iloc[i, 2], track.iloc[i, 3], track.iloc[i, 4], track.iloc[i, 5]]
                 tracks.at[index, 6] = StopAnalyzer.iou(bb0, bb1)
 
-                center = Point(track.iloc[i,2]+track.iloc[i,4], track.iloc[i,3]+track.iloc[i,5])
-                for j in range(len(self.vzones)):
-                    if center.within(self.vzones[j]):
-                        tracks.loc[tracks[1]==vehicle, 7]=j
-                        break
+                if len(self.vzones)>0:
+                    center = Point(track.iloc[i,2]+track.iloc[i,4], track.iloc[i,3]+track.iloc[i,5])
+                    for j in range(len(self.vzones)):
+                        if center.within(self.vzones[j]):
+                            tracks.loc[tracks[1]==vehicle, 7]=j
+                            break
                 
-                bb = box(track.iat[i, 2], track.iat[i, 3], track.iat[i, 2] + track.iat[i, 4], track.iat[i, 3] + track.iat[i, 5])
-                for j in range(len(self.hzones)):
-                    if bb.intersects(self.hzones[j]):
-                        if j > tracks.at[index, 8]:
-                            tracks.at[index, 8] = j
+                if len(self.hzones)>0:
+                    bb = box(track.iat[i, 2], track.iat[i, 3], track.iat[i, 2] + track.iat[i, 4], track.iat[i, 3] + track.iat[i, 5])
+                    for j in range(len(self.hzones)):
+                        if bb.intersects(self.hzones[j]):
+                                pass
+                            #if j > tracks.at[index, 8]:
+                                #tracks.at[index, 8] = j
+                                
+                                
             if self.verbose:
                 pbar.update()
 
@@ -64,13 +81,13 @@ class StopAnalyzer():
 
         return tracks
     
-    def __event_identify(self, tracks, video_index, video_tot):
+    def event_identify(self, tracks:pd.DataFrame, video_index:int=None, video_tot:int=None)->pd.DataFrame:
 
         pbar = tqdm(total=len(tracks), unit=' frames')
         if video_index and video_tot:
-            pbar.set_description_str("Identify events {} of {}".format(video_index, video_tot))
+            pbar.set_description_str("Identifying events {} of {}".format(video_index, video_tot))
         else:
-            pbar.set_description_str("Identify events ")
+            pbar.set_description_str("Identifying events ")
 
         for i in range(len(tracks)):
 
@@ -100,8 +117,7 @@ class StopAnalyzer():
 
         return tracks
     
-    def __event_count(self, tracks, video_index, video_tot):
-
+    def event_count(self, tracks:pd.DataFrame, video_index:int=None, video_tot:int=None):
 
         pbar = tqdm(unit='events')
         results = []
@@ -111,9 +127,9 @@ class StopAnalyzer():
             vehicles = tracks.loc[tracks[9]==key][1].unique()
 
             if video_index and video_tot:
-                pbar.set_description_str("Count event {} for {} of {}".format(key, video_index, video_tot))
+                pbar.set_description_str("Counting event {} for {} of {}".format(key, video_index, video_tot))
             else:
-                pbar.set_description_str("Count event {}".format(key))
+                pbar.set_description_str("Counting event {}".format(key))
 
             pbar.total = len(vehicles)
             for vehicle in vehicles:
@@ -133,13 +149,55 @@ class StopAnalyzer():
         return results
     
     @staticmethod
-    def export_label(track_file, analysis_file, label_file, vid_field=1, label_field=10, frame_field=0,
-                     event_label=None, vid_disp=True, verbose=True):
+    def add_timestamp(result_file:str, 
+                      output_file:str,
+                      init_time:datetime.datetime, 
+                      fps:float,
+                      verbose:bool=True) -> pd.DataFrame:
+        
+        results = pd.read_csv(result_file)
+        for index, row in results.iterrows():
+            start_frame = row['START_FRAME']
+            end_frame = row['END_FRAME']
+            diff_frame = end_frame -  start_frame
+
+            start_time = init_time + datetime.timedelta(seconds=start_frame/fps)
+            duration = diff_frame/fps
+
+            results.at[index, 'YEAR'] = start_time.year
+            results.at[index, 'MONTH'] = start_time.month
+            results.at[index, 'HOUR'] = start_time.hour
+            results.at[index, 'MINUTE'] = start_time.minute
+            results.at[index, 'SECOND'] = start_time.second
+            results.at[index, 'DURATION'] = duration
+            results.at[index, 'FPS'] = fps
+            results.at[index, 'START_SEC'] = start_frame/fps
+    
+        results.to_csv(output_file, index=False)
+        if verbose:
+            print('Writen to {}'.format(output_file))
+
+    @staticmethod
+    def export_label(track_file:str, 
+                     analysis_file:str, 
+                     label_file:str, 
+                     vid_field:int=1, 
+                     label_field:int=10, 
+                     frame_field:int=0,
+                     event_label:list[str]=None, 
+                     vid_disp:bool=True, 
+                     verbose:bool=True, 
+                     video_index:int=None, video_tot:int=None):
         
         tracks = pd.read_csv(track_file, header=None)
         results = pd.read_csv(analysis_file)
 
-        pbar = tqdm(total=len(tracks), desc='Generate labels ', unit=' frames')
+        pbar = tqdm(total=len(tracks), unit=' frames')
+        if video_index and video_tot:
+            pbar.set_description_str("Generating labels {} of {}".format(video_index, video_tot))
+        else:
+            pbar.set_description_str("Generating labels ")
+
         for index, track in tracks.iterrows():
 
             vid = int(track[vid_field])
@@ -147,9 +205,10 @@ class StopAnalyzer():
             selected = results.loc[(results['START_FRAME']<=track[frame_field]) & 
                                    (results['END_FRAME']>=track[frame_field]) &
                                    (results['TRACKID']==track[vid_field])]
+            
             if len(selected)>0:
-                
-                event = selected.iloc[0]['EVENT']
+
+                event = int(selected.iloc[0]['EVENT'])
 
                 if event_label:
                     tracks.at[index, label_field] = str(vid)+"-"+str(event_label[event])
@@ -174,7 +233,7 @@ class StopAnalyzer():
 
 
     @staticmethod
-    def iou(bb1, bb2):
+    def iou(bb1:tuple[int, int, int, int], bb2:tuple[int, int, int, int]):
         """
         Calculate the Intersection over Union (IoU) of two bounding boxes.
         Parameters
@@ -221,7 +280,7 @@ class StopAnalyzer():
         return iou
     
     @staticmethod
-    def gen_zones(line_coords):
+    def gen_zones(line_coords:list[list[tuple[int, int], tuple[int, int]]])->list[Polygon]:
         """
             Generate a list of shapely polygons
             Inputs:

dnt/analysis/stop2.py

@@ -0,0 +1,289 @@
+from shapely.geometry import Point, Polygon, LineString, box
+import geopandas as gpd, pandas as pd
+import datetime
+from tqdm import tqdm
+
+class StopAnalyzer():
+    def __init__(self, 
+                 hzones: list[Polygon] = [], 
+                 vzones: list[Polygon] = [],
+                 event_dict:dict[int, list[int]] = {}, 
+                 stop_iou:float=0.97, 
+                 frame_buffer:int=5,
+                 verbose:bool=True):
+        
+        self.hzones = hzones
+        self.vzones = vzones
+        self.event_dict = event_dict
+        self.stop_iou = stop_iou
+        self.frame_buffer = frame_buffer
+        self.verbose = verbose
+
+    def analyze(self, tracks:pd.DataFrame=None, track_file:str=None,video_index:int=None, video_tot:int=None)->list[pd.DataFrame, pd.DataFrame]:
+
+        if tracks is None:
+            tracks = pd.read_csv(track_file, header=None)
+        tracks = StopAnalyzer.add_field_names(tracks)
+
+        tracks = self.scan_stop(tracks, video_index, video_tot)
+        tracks = self.identify_event(tracks, video_index, video_tot)
+        events = self.count_event(tracks, video_index, video_tot)
+
+        return tracks, events
+    
+    def scan_stop(self, tracks:pd.DataFrame, video_index:int=None, video_tot:int=None)->pd.DataFrame:
+        
+        ids = tracks['track'].unique()
+        tracks['hzone'] = -1
+        
+        pbar = tqdm(total=len(ids), unit=' tracks')
+        if video_index and video_tot:
+            pbar.set_description_str("Scanning stops {} of {}".format(video_index, video_tot))
+        else:
+            pbar.set_description_str("Scanning stops")
+
+        for id in ids:
+            track = tracks[tracks['track'] == id].sort_values(by='frame') # sort by frame in ascending
+
+            for i in range(self.frame_buffer, len(track)):
+                index = track.iloc[i].name
+                bb0 = [track.iloc[i-self.frame_buffer]['x'], track.iloc[i-self.frame_buffer]['y'], 
+                       track.iloc[i-self.frame_buffer]['w'], track.iloc[i-self.frame_buffer]['h']]
+                bb1 = [track.iloc[i]['x'], track.iloc[i]['y'], track.iloc[i]['w'], track.iloc[i]['h']]
+                iou_score = StopAnalyzer.iou(bb0, bb1)
+
+                tracks.at[index, 'iou'] = iou_score 
+                if iou_score >= self.stop_iou:
+                    tracks.at[index, 'stop'] = 1
+
+                if len(self.vzones)>0:
+                    center = Point(track.iloc[i]['x']+track.iloc[i]['w'], track.iloc[i]['y']+track.iloc[i]['h'])
+                    for j in range(len(self.vzones)):
+                        if center.within(self.vzones[j]):
+                            tracks.loc[tracks['track']==id, 'vzone']=j
+                            break
+
+                if len(self.hzones)>0:
+                    bb = box(track.iloc[i]['x'], track.iloc[i]['y'], track.iloc[i]['x'] 
+                             + track.iloc[i]['w'], track.iloc[i]['y'] + track.iloc[i]['h'])
+                    for j in range(len(self.hzones)):
+                        if bb.intersects(self.hzones[j]):
+                            if j > tracks.at[index, 'hzone']:
+                                tracks.at[index, 'hzone'] = j
+                                                 
+            if self.verbose:
+                pbar.update()
+
+        pbar.close()
+        return tracks  
+
+    def identify_event(self, tracks:pd.DataFrame, video_index:int=None, video_tot:int=None)->pd.DataFrame:
+
+        pbar = tqdm(total=len(tracks), unit=' frames')
+        if video_index and video_tot:
+            pbar.set_description_str("Identifying events {} of {}".format(video_index, video_tot))
+        else:
+            pbar.set_description_str("Identifying events")
+
+        for i in range(len(tracks)):
+
+            if tracks.iloc[i]['stop'] == 1:
+                vzone = tracks.iloc[i]['vzone']
+                hzone = tracks.iloc[i]['hzone']
+                frame = tracks.iloc[i]['frame']
+                id = tracks.iloc[i]['track']
+                vehs_inlane = tracks.loc[(tracks['frame']==frame) & (tracks['track']!=id) & (tracks['vzone']==vzone) & (tracks['hzone']>hzone)]
+
+                if len(vehs_inlane)==0:
+                    pre_key = -1
+                    for key in self.event_dict:
+                        if key >= pre_key:
+                            pre_key = key
+                            if tracks.iloc[i]['hzone'] in self.event_dict[key]:
+                                tracks.at[i, 'event'] = key
+
+            if self.verbose:
+                pbar.update()
+        
+        pbar.close()
+
+        return tracks
+
+    def count_event(self, tracks:pd.DataFrame, video_index:int=None, video_tot:int=None)->pd.DataFrame:
+
+        pbar = tqdm(unit='events')
+        results = []
+        for key in self.event_dict:
+
+            vehicles = tracks.loc[tracks['event']==key]['track'].unique()
+
+            if video_index and video_tot:
+                pbar.set_description_str("Counting event {} for {} of {}".format(key, video_index, video_tot))
+            else:
+                pbar.set_description_str("Counting event {}".format(key))
+
+            pbar.total = len(vehicles)
+            for vehicle in vehicles:
+
+                track = tracks[(tracks['track'] == vehicle) & (tracks['event'] == key)]
+
+                start_frame = int(track['frame'].min())
+                end_frame = int(track['frame'].max())
+                vzone = track['vzone'].mode()[0]
+                results.append([key, vehicle, vzone, start_frame, end_frame])
+                
+                if self.verbose:
+                    pbar.update()
+
+        results = pd.DataFrame(results, columns=['event', 'track', 'vzone', 'start_frame', 'end_frame'])
+        pbar.close()
+
+        return results
+    
+    def generate_labels(self, tracks:pd.DataFrame, events:pd.DataFrame, video_index:int=None, video_tot:int=None)->pd.DataFrame:
+        
+        event_tracks = tracks.loc[tracks['event']>-1].copy()
+        for index, event in events.iterrows():
+            track_id = event['track']
+            start_frame = event['start_frame']
+            end_frame = event['end_frame']
+            event_id = event['event']
+
+            event_tracks.loc[(event_tracks['track']==track_id) & (event_tracks['frame']>=start_frame) 
+                             & (event_tracks['frame']<=end_frame), 'event'] = event_id
+
+        pbar = tqdm(total= len(event_tracks))        
+        if video_index and video_tot:
+            pbar.set_description_str("Generating labeles for {} of {}".format(video_index, video_tot))
+        else:
+            pbar.set_description_str("Generating labels")
+
+        results = []
+        for index, track in event_tracks.iterrows():
+            desc = 'Stop'
+            if track['event'] == 1:
+                desc = 'Pass'
+            elif track['event'] == 2:
+                desc = 'Invading'
+
+            results.append([track['frame'], 'bbox', [(track['x'], track['y']), (track['x']+track['w'], track['y']+track['h'])], 
+                            (0, 255,0), 1, 1, desc])
+            
+            if self.verbose:
+                pbar.update()
+
+        df = pd.DataFrame(results, columns=['frame','type','coords','color','size','thick','desc'])
+        df.sort_values(by='frame')
+        
+        return df
+
+    
+    @staticmethod
+    def add_field_names(tracks: pd.DataFrame)->pd.DataFrame:
+        if len(tracks.columns)!=10:
+            raise Exception('The number of fields is not nine.')
+        tracks.columns = ['frame', 'track', 'x', 'y', 'w', 'h', 'score', 'cls', 'r3','r4']
+        tracks['iou'] = -1.0
+        tracks['stop'] = -1
+        tracks['vzone'] = -1
+        tracks['hzone'] = -1
+        tracks['event'] = -1   
+        return tracks
+    
+    @staticmethod
+    def iou(bb1:tuple[int, int, int, int], bb2:tuple[int, int, int, int]):
+        """
+        Calculate the Intersection over Union (IoU) of two bounding boxes.
+        Parameters
+        ----------
+        bb1 : [x1, y1, w, h]
+            x1, y1 - top left corner
+            w, h - width and height
+        bb2 : [x1, y1, w, h]
+            x1, y1 - top left corner
+            w, h - width and height
+        Returns
+        -------
+        iou: float [0, 1]
+        """
+
+        assert bb1[0] < bb1[0] + bb1[2]
+        assert bb1[1] < bb1[1] + bb1[3]
+        assert bb2[0] < bb2[0] + bb2[2]
+        assert bb2[1] < bb2[1] + bb2[3]
+
+        # determine the coordinates of the intersection rectangle
+        x_left = max(bb1[0], bb2[0])
+        y_top = max(bb1[1], bb2[1])
+        x_right = min(bb1[0]+bb1[2], bb2[0]+bb2[2])
+        y_bottom = min(bb1[1]+bb1[3], bb2[1]+bb2[3])
+
+        if x_right < x_left or y_bottom < y_top:
+            return 0.0
+
+        # The intersection of two axis-aligned bounding boxes is always an
+        # axis-aligned bounding box
+        intersection_area = (x_right - x_left) * (y_bottom - y_top)
+
+        # compute the area of both AABBs
+        bb1_area = bb1[2] * bb1[3]
+        bb2_area = bb2[2] * bb2[3]
+
+        # compute the intersection over union by taking the intersection
+        # area and dividing it by the sum of prediction + ground-truth
+        # areas - the interesection area
+        iou = intersection_area / float(bb1_area + bb2_area - intersection_area)
+        assert iou >= 0.0
+        assert iou <= 1.0
+        return iou
+
+    @staticmethod
+    def gen_zones(line_coords:list[list[tuple[int, int], tuple[int, int]]])->list[Polygon]:
+        """
+            Generate a list of shapely polygons
+            Inputs:
+                line_coords: a list of line coords ([[(x11, y11),(x12, y12)], [(x21, y21),(x22, y22)], ...])
+            Returns:
+                A list of PloyGons  
+        """
+
+        zones = []
+        for i in range(1, len(line_coords)):
+            shell = line_coords[i-1] + [line_coords[i][1], line_coords[i][0]]
+            zones.append(Polygon(shell))        
+
+        return zones
+      
+if __name__=='__main__':
+    track_file = '/mnt/d/videos/ped2stage/tracks/gh021291_track_veh.txt'
+    tracks = pd.read_csv(track_file, header=None)
+    tracks['r1'] = -1
+    tracks['r2'] = -1
+    tracks['r3'] = -1
+    tracks = StopAnalyzer.add_field_names(tracks)
+
+    v_coords = []
+    h_coords = [[(194, 769),(480, 484)], 
+          [(550, 812),(767, 487)],
+          [(677, 835),(824, 484)],
+          [(1087, 870),(958, 483)]]
+    hzones = StopAnalyzer.gen_zones(h_coords)
+
+    event_dict = {
+        0: [0],    # Stop before stop bar
+        1: [1],    # Pass stop bar
+        2: [2]     # Invading crosswalk
+        }
+    analyzer = StopAnalyzer(hzones=hzones, event_dict=event_dict)
+
+    tracks_stop = analyzer.scan_stop(tracks)
+    tracks_event = analyzer.identify_event(tracks_stop)
+    tracks_event.to_csv('/mnt/d/videos/ped2stage/stop/gh021291_track_veh_stop.txt', index=False)
+
+    event_count = analyzer.count_event(tracks_event)
+    event_count.to_csv('/mnt/d/videos/ped2stage/stop/gh021291_track_stops_count.txt', index=False)
+
+    labels = analyzer.generate_labels(tracks_event, event_count)
+    labels.to_csv('/mnt/d/videos/ped2stage/stop/gh021291_track_stops_label.txt', index=False)
+
+