Create tracker2.py

tracker2.py

@@ -0,0 +1,129 @@
+import cv2
+import math
+import time
+import numpy as np
+
+limit = 80 #km/hr
+
+file = open("D://TrafficRecord//SpeedRecord.txt","w")
+file.write("ID \t SPEED\n------\t-------\n")
+file.close()
+
+
+class EuclideanDistTracker:
+    def __init__(self):
+        # Store the center positions of the objects
+        self.center_points = {}
+
+        self.id_count = 0
+        #self.start = 0
+        #self.stop = 0
+        self.et=0
+        self.s1 = np.zeros((1,1000))
+        self.s2 = np.zeros((1,1000))
+        self.s = np.zeros((1,1000))
+        self.f = np.zeros(1000)
+        self.capf = np.zeros(1000)
+        self.count = 0
+        self.exceeded = 0
+
+
+    def update(self, objects_rect):
+        objects_bbs_ids = []
+
+        # Get center point of new object
+        for rect in objects_rect:
+            x, y, w, h = rect
+            cx = (x + x + w) // 2
+            cy = (y + y + h) // 2
+
+            #CHECK IF OBJECT IS DETECTED ALREADY
+            same_object_detected = False
+
+            for id, pt in self.center_points.items():
+                dist = math.hypot(cx - pt[0], cy - pt[1])
+
+                if dist < 70:
+                    self.center_points[id] = (cx, cy)
+                    objects_bbs_ids.append([x, y, w, h, id])
+                    same_object_detected = True
+
+                    #START TIMER
+                    if (y >= 410 and y <= 430):
+                        self.s1[0,id] = time.time()
+
+                    #STOP TIMER and FIND DIFFERENCE
+                    if (y >= 235 and y <= 255):
+                        self.s2[0,id] = time.time()
+                        self.s[0,id] = self.s2[0,id] - self.s1[0,id]
+
+                    #CAPTURE FLAG
+                    if (y<235):
+                        self.f[id]=1
+
+
+            #NEW OBJECT DETECTION
+            if same_object_detected is False:
+                self.center_points[self.id_count] = (cx, cy)
+                objects_bbs_ids.append([x, y, w, h, self.id_count])
+                self.id_count += 1
+                self.s[0,self.id_count]=0
+                self.s1[0,self.id_count]=0
+                self.s2[0,self.id_count]=0
+
+        # ASSIGN NEW ID to OBJECT
+        new_center_points = {}
+        for obj_bb_id in objects_bbs_ids:
+            _, _, _, _, object_id = obj_bb_id
+            center = self.center_points[object_id]
+            new_center_points[object_id] = center
+
+        self.center_points = new_center_points.copy()
+        return objects_bbs_ids
+
+    #SPEEED FUNCTION
+    def getsp(self,id):
+        if (self.s[0,id]!=0):
+            s = 214.15 / self.s[0, id]
+        else:
+            s = 0
+
+        return int(s)
+
+    #SAVE VEHICLE DATA
+    def capture(self,img,x,y,h,w,sp,id):
+        if(self.capf[id]==0):
+            self.capf[id] = 1
+            self.f[id]=0
+            crop_img = img[y-5:y + h+5, x-5:x + w+5]
+            n = str(id)+"_speed_"+str(sp)
+            file = 'D://TrafficRecord//' + n + '.jpg'
+            cv2.imwrite(file, crop_img)
+            self.count += 1
+            filet = open("D://TrafficRecord//SpeedRecord.txt", "a")
+            if(sp>limit):
+                file2 = 'D://TrafficRecord//exceeded//' + n + '.jpg'
+                cv2.imwrite(file2, crop_img)
+                filet.write(str(id)+" \t "+str(sp)+"<---exceeded\n")
+                self.exceeded+=1
+            else:
+                filet.write(str(id) + " \t " + str(sp) + "\n")
+            filet.close()
+
+
+    #SPEED_LIMIT
+    def limit(self):
+        return limit
+
+    #TEXT FILE SUMMARY
+    def end(self):
+        file = open("D://TrafficRecord//SpeedRecord.txt", "a")
+        file.write("\n-------------\n")
+        file.write("-------------\n")
+        file.write("SUMMARY\n")
+        file.write("-------------\n")
+        file.write("Total Vehicles :\t"+str(self.count)+"\n")
+        file.write("Exceeded speed limit :\t"+str(self.exceeded))
+        file.close()
+
+