Merge branch 'nr' of https://github.com/RizwanMunawar/streamgrid into nr

Author: RizwanMunawar

streamgrid/grid.py

@@ -15,7 +15,9 @@
 class StreamGrid:
     """Ultra-fast multi-stream video display with object detection."""
 
-    def __init__(self, sources=None, model=None, save=True, device="cpu", analytics=False):
+    def __init__(
+        self, sources=None, model=None, save=True, device="cpu", analytics=False
+    ):
         # Initialize components
         self.stream_manager = StreamManager(sources)
         self.model = model
@@ -32,7 +34,9 @@ def __init__(self, sources=None, model=None, save=True, device="cpu", analytics=
         self.plotter = StreamAnnotator(self.cell_w, self.cell_h, self.max_sources)
 
         # Display state
-        self.grid = np.zeros((self.rows * self.cell_h, self.cols * self.cell_w, 3), dtype=np.uint8)
+        self.grid = np.zeros(
+            (self.rows * self.cell_h, self.cols * self.cell_w, 3), dtype=np.uint8
+        )
         self.frames = {}
         self.show_stats = True
         self.running = False
@@ -51,8 +55,10 @@ def setup_video_writer(self):
         """Setup video writer for saving output."""
         fourcc = cv2.VideoWriter_fourcc(*"mp4v")
         return cv2.VideoWriter(
-            f"streamgrid_output_{self.max_sources}_streams.mp4", fourcc, 30,
-            (self.cols * self.cell_w, self.rows * self.cell_h)
+            f"streamgrid_output_{self.max_sources}_streams.mp4",
+            fourcc,
+            30,
+            (self.cols * self.cell_w, self.rows * self.cell_h),
         )
 
     def process_batch(self):
@@ -79,7 +85,9 @@ def process_batch(self):
 
             # Run inference if model available
             if self.model:
-                results = self.model.predict(frames, conf=0.25, verbose=False, device=self.device)
+                results = self.model.predict(
+                    frames, conf=0.25, verbose=False, device=self.device
+                )
                 for source_id, frame, result in zip(ids, frames, results):
                     self.update_source(source_id, frame, result)
             else:
@@ -110,10 +118,14 @@ def update_source(self, source_id, frame, results=None):
             detections = 0
             if results and results.boxes is not None:
                 detections = len(results.boxes)
-                resized = self.plotter.draw_detections(resized, results, frame.shape[:2])
+                resized = self.plotter.draw_detections(
+                    resized, results, frame.shape[:2]
+                )
 
             # Add source label
-            resized = self.plotter.draw_source_label(resized, source_id, self.show_stats)
+            resized = self.plotter.draw_source_label(
+                resized, source_id, self.show_stats
+            )
 
             # Store processed frame
             self.frames[source_id] = resized
@@ -140,8 +152,10 @@ def update_display(self):
         # Add FPS overlay
         if self.show_stats:
             self.grid = self.plotter.draw_fps_overlay(
-                self.grid, self.prediction_fps,
-                self.cols * self.cell_w, self.rows * self.cell_h
+                self.grid,
+                self.prediction_fps,
+                self.cols * self.cell_w,
+                self.rows * self.cell_h,
             )
 
         # Display and save
@@ -166,7 +180,7 @@ def run(self):
                 key = cv2.waitKey(1) & 0xFF
                 if key == 27:  # ESC
                     break
-                elif key == ord('s'):
+                elif key == ord("s"):
                     self.show_stats = not self.show_stats
         finally:
             self.stop()
@@ -185,7 +199,9 @@ def stop(self):
             self.analytics.summary()
         if self.video_writer:
             self.video_writer.release()
-            LOGGER.info(f"✅ Video saved: streamgrid_output_{self.max_sources}_streams.mp4")
+            LOGGER.info(
+                f"✅ Video saved: streamgrid_output_{self.max_sources}_streams.mp4"
+            )
 
         with self.lock:
             self.frames.clear()