Update app.py

app.py

@@ -21,7 +21,7 @@ db = firestore.client()
 colors = [(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) 
           for j in range(10)]
 
-detection_threshold = 0.7
+detection_threshold = 0.1
 model = YOLO(MODEL_PATH)
 
 def addToDatabase(ss_id, obj_ids):
@@ -48,7 +48,7 @@ def addToDatabase(ss_id, obj_ids):
 
 
 def traffic_counting(video):
-
+   
     obj_ids = {"person": [], 
                   "bicycle": [], 
                   "car": [], 
@@ -97,25 +97,72 @@ def traffic_counting(video):
         # Count each type of traffic
         output_data = {key: len(value) for key, value in obj_ids.items()}
         df = pd.DataFrame(list(output_data.items()), columns=['Type', 'Number'])
-
+       
         yield frame, df
         ret, frame = cap.read()
-
+        
+    
     cap.release()
     cv2.destroyAllWindows()
     video_path = video.replace("\\", "/")
-    addToDatabase(video_path.split("/")[-1][:-4], obj_ids)
+    # addToDatabase(video_path.split("/")[-1][:-4], obj_ids)
 
 
+# input_video = gr.Video(label="Input Video")
+# output_video = gr.outputs.Video(label="Processing Video")
+# output_data = gr.Dataframe(interactive=False, label="Traffic's Frequency")
 
-input_video = gr.Video(label="Input Video")
-output_video = gr.Image(type="numpy", label="Processing Video")
-output_data = gr.Dataframe(interactive=False, label="Traffic's Frequency")
+# demo = gr.Interface(traffic_counting,
+#                     inputs=input_video,
+#                     outputs=[output_video, output_data],
+#                     examples=[os.path.join('video', x) for x in os.listdir('video') if x != ".gitkeep"],
+#                     allow_flagging='never'
+#                     )
+def traffic_detection(image):
+    
+    results = model.predict(image)
+    detections = []
+    obj_ids = {"person": [], 
+                  "bicycle": [], 
+                  "car": [], 
+                  "motocycle": [], 
+                  "bus": [], 
+                  "truck": [], 
+                  "other": []}
 
-demo = gr.Interface(traffic_counting,
-                    inputs=input_video,
-                    outputs=[output_video, output_data],
-                    examples=[os.path.join('video', x) for x in os.listdir('video') if x != ".gitkeep"],
+    for result in results:
+        for r in result.boxes.data.tolist():
+            x1, y1, x2, y2, score, class_id = r
+            x1 = int(x1)
+            x2 = int(x2)
+            y1 = int(y1)
+            y2 = int(y2)
+            class_id = int(class_id)
+            if score > detection_threshold:
+                detections.append([x1, y1, x2, y2, class_id, score])
+            cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (compute_color_for_labels(class_id)), 1)
+            label_name = ID2LABEL[class_id] if class_id in ID2LABEL.keys() else "other"
+            cv2.putText(image,f"{label_name}", 
+                            (int(x1) + 5, int(y1) - 5), 
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.3,compute_color_for_labels(class_id), 1, cv2.LINE_AA )         
+       
+        # Count each type of traffic
+        output_data = {key: len(value) for key, value in obj_ids.items()}
+        df = pd.DataFrame(list(output_data.items()), columns=['Type', 'Number'])
+        yield image, df
+
+        
+                
+    
+
+# Input is a image
+input_image = gr.Image(label="Input Image")
+output_image = gr.Image(type="filepath", label="Processing Image")
+output_data = gr.Dataframe(interactive=False, label="Traffic's Frequency")
+demo = gr.Interface(traffic_detection,
+                    inputs=input_image,
+                    outputs=[output_image, output_data],
+                    examples=[os.path.join('image', x) for x in os.listdir('image') if x != ".gitkeep"],
                     allow_flagging='never'
                     )