Merge remote-tracking branch 'origin/add-capture-card' into add-capture-card

# Conflicts:
#	core.py
#	server.py

app.py

@@ -4,9 +4,12 @@ import sys
 from datetime import datetime, timedelta
 from time import sleep
 from typing import List
-
-import sys
 import os
+import time
+import threading  # Import threading module
+
+from cvStreamer import cvStreamer
+from detector import Detector
 
 # Add the proto directory to the Python path
 proto_dir = os.path.join(os.path.dirname(__file__), 'message_queue', 'proto')
@@ -24,6 +27,7 @@ from icecream import ic
 
 from configs import ConfigManager
 from core import Core
+from tracker import Tracker
 from message_queue.Bridge import Bridge
 from message_queue.Manager import Manager
 from message_queue.proto.ImageMessage_pb2 import ImageMessage, TrackMode
@@ -37,13 +41,20 @@ from message_queue.proto.ConnectStatus_pb2 import ConnectStatus
 from video_streamer.gst_video_streamer import GstVideoStreamer
 import cv2
 
-import os
-import time  # Ensure time module is imported
 
 config_manager = ConfigManager('config.yaml')
 rtsp_links = config_manager.configs['rtsp_links'].get()
 debug = config_manager.configs['debug'].get()
 
+def findUsbCams():
+    cv_cameras = []
+    for i in range(50):
+        cap = cvStreamer(i)
+        if cap.isOpened():
+            cv_cameras.append(cap)
+    return cv_cameras
+
+
 def handle_camera_status(status: int):
     m = Message()
     m.msgType = MessageType.MESSAGE_TYPE_CAMERA_CONNECTION_STATUS
@@ -53,72 +64,145 @@ def handle_camera_status(status: int):
 # Helper class to track client connection status
 class ConnectionTracker:
     def __init__(self):
-        self.last_message_time = time.time()  # Use time.time() to get the current time
-        self.timeout = 15  # Timeout in seconds (adjust as needed)
+        self.last_message_time = time.time()
+        self.timeout = 15  # 15-second timeout
+        self.last_client_ip = None  # Track last connected client IP
+
+    def update_last_message_time(self, client_ip=None):
+        self.last_message_time = time.time()
+        if client_ip:
+            self.last_client_ip = client_ip  # Update last known client
+
+    def is_client_active(self):
+        """Check if the last client is still active within the last 15 seconds."""
+        return (time.time() - self.last_message_time) < self.timeout
+# Create an instance of ConnectionTracker
+connection_tracker = ConnectionTracker()
 
-    def update_last_message_time(self):
-        self.last_message_time = time.time()  # Update with the current time
 
-    def is_client_connected(self):
-        return (time.time() - self.last_message_time) < self.timeout  # Use time.time()
+class ConnectionThread(threading.Thread):
+    def __init__(self):
+        super().__init__()
+        self.running = True
 
-# Create an instance of ConnectionTracker
-connection_tracker = ConnectionTracker()
+    def run(self):
+        while self.running:
 
-def check_client_connection():
-    if not connection_tracker.is_client_connected():
-        print("Client disconnected. Searching for another client...")
-        cl_ip, _ = start_discovery_service(12345)
-        ic(cl_ip)
 
-        # Reinitialize the Manager with the new client addressimport time  # Ensure time module is imported
 
-# Helper class to track client connection status
-class ConnectionTracker:
-    def __init__(self):
-        self.last_message_time = time.time()  # Use time.time() to get the current time
-        self.timeout = 15  # Timeout in seconds (adjust as needed)
+            sleep(0.5)
 
-    def update_last_message_time(self):
-        self.last_message_time = time.time()  # Update with the current time
+            if connection_tracker.last_client_ip:
 
-    def is_client_connected(self):
-        return (time.time() - self.last_message_time) < self.timeout  # Use time.time()
+                if debug:
+                    print(f"Checking if last client {connection_tracker.last_client_ip} is still available...")
 
-# Create an instance of ConnectionTracker
-connection_tracker = ConnectionTracker()
+                if self.query_last_client(connection_tracker.last_client_ip):
+                    connection_tracker.update_last_message_time()
+
+                    if debug:
+                        print(f"Last client {connection_tracker.last_client_ip} responded. Continuing...")
+                    continue  # Skip discovering a new client
+            if not connection_tracker.is_client_active():
+
+
+                print("Client inactive for 15 seconds. Searching for another client...")
+                cl_ip, _ = self.start_discovery_service(12345)
+                print(f"New client found: {cl_ip}")
+
+                # Reinitialize the Manager with the new client address
+                global manager
+                manager = Manager(f"tcp://{cl_ip}:5558", f"tcp://{cl_ip}:5557")
+                manager.start(manager_callback)
+
+                connection_tracker.update_last_message_time(cl_ip)  # Update with new client
+
+    def query_last_client(self, client_ip):
+        """Send a heartbeat packet to the last known client IP on port 29170 and wait for a response."""
+
+        if debug:
+            print(f"Sending heartbeat to {client_ip}:29170")  # Debugging
+
+        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        sock.settimeout(5)  # 5-second timeout
+        heartbeat_port = 29170  # New port for heartbeat
+
+        try:
+            sock.sendto(b"HEARTBEAT", (client_ip, heartbeat_port))  # Send heartbeat message
+            data, addr = sock.recvfrom(1024)  # Wait for response
+
+            if debug:
+                print(f"Received response from {addr}: {data.decode()}")  # Debugging
+
+            if data.decode() == "HEARTBEAT_ACK":
+                connection_tracker.update_last_message_time()  # Update the last message time
 
-def check_client_connection():
-    if not connection_tracker.is_client_connected():
-        print("Client disconnected. Searching for another client...")
-        cl_ip, _ = start_discovery_service(12345)
-        print(cl_ip)
+                if debug:
+                    print(f"Client {client_ip} is still responding.")
+                return True
 
-        # Reinitialize the Manager with the new client address
-        global manager
-        manager = Manager(f"tcp://{cl_ip}:5558", f"tcp://{cl_ip}:5557")
-        manager.start(manager_callback)  # Restart the Manager and re-register the callback
+        except socket.timeout:
+
+            print(f"Client {client_ip} did not respond. Marking as inactive.")
+
+        finally:
+            sock.close()
+
+        return False  # Client did not respond
+
+
+    def start_discovery_service(self, port):
+        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        sock.bind(('0.0.0.0', port))
+        print(f"Discovery service listening on port {port}...")
+
+        while True:
+            data, addr = sock.recvfrom(1024)
+            if data.decode() == "DISCOVER_SERVER":
+                print(f"Received discovery request from {addr}")
+                sock.sendto(b"SERVER_RESPONSE", addr)
+                break
+        sock.close()
+        return addr
+
+
+
+
+    def stop(self):
+        self.running = False
 
-        connection_tracker.update_last_message_time()  # Reset the timer after reconnecting
 
 if __name__ == '__main__':
     QCoreApplication.setAttribute(Qt.AA_EnableHighDpiScaling)
     app = QCoreApplication(sys.argv)
 
-    videoStreamers: List[GstVideoStreamer] = []
-    streamerThreads = []
+        # cv_cameras = []
+        # for i in range(3):
+        #     cap = cvStreamer(i)
+        #     if cap.isOpened():
+        #         cv_cameras.append(cap)
+
+
+    videoStreamers = []
     for idx, rtsp_link in enumerate(rtsp_links):
         videoStreamer = GstVideoStreamer(rtsp_link, [1920, 1080, 3], str(idx), fps=15)
         videoStreamer.cameraStatus.connect(handle_camera_status)
         print(f'{videoStreamer.id} connected')
         videoStreamers.append(videoStreamer)
 
-    core = Core(videoStreamers)
+    tracker = Tracker()
+    detector = Detector(classes=[0, 2, 5, 7])
+
+    sources = findUsbCams()
+    if len(sources) >= 2:
+        core = Core(sources,tracker,detector)
+    else:
+        core = Core(videoStreamers, tracker, detector)
+
 
     def manager_callback(msg_str):
         msg = Message()
         msg.ParseFromString(msg_str)
-        connection_tracker.update_last_message_time()  # Update the last message time
 
         if msg.msgType == MessageType.MESSAGE_TYPE_TOGGLE_TRACK:
             if msg.track_cmd.command == TrackCommand.TRACK_COMMAND_STOP:
@@ -130,18 +214,34 @@ if __name__ == '__main__':
         elif msg.msgType == MessageType.MESSAGE_TYPE_TRACK_SETTINGS:
             core.set_thickness(msg.track_settings.thickness)
         elif msg.msgType == MessageType.MESSAGE_TYPE_SWITCH_CAMERA:
+            print(f"switch camera detected ,primary value is : {msg.cam_switch.primaryCamType}")
+            print(f"switch camera detected ,secondary value is : {msg.cam_switch.secondaryCamType}")
             core.set_source(msg.cam_switch.primaryCamType)
         elif msg.msgType == MessageType.MESSAGE_TYPE_SET_CAMERA:
-            ip = msg.cam_set.ip
-            videoStreamers = []
-            for src in range(2):
-                rtsp_link = f'rtsp://admin:Abc.12345@{ip}:554/ch{src}/stream0'
-                # ic(rtsp_link)
-                videoStreamer = GstVideoStreamer(rtsp_link, [1920, 1080, 3], str(idx), fps=15)
-                videoStreamer.cameraStatus.connect(handle_camera_status)
-                print(f'{videoStreamer.id} connected')
-                videoStreamers.append(videoStreamer)
-            core.set_video_sources(videoStreamers)
+            if msg.cam_set.cameraSource == 1:
+                print(f"set camera source to network")
+                ip = msg.cam_set.ip
+                videoStreamers = []
+                for src in range(2):
+                    rtsp_link = f'rtsp://admin:Abc.12345@{ip}:554/ch{src}/stream0'
+                    # ic(rtsp_link)
+                    videoStreamer = GstVideoStreamer(rtsp_link, [1920, 1080, 3], str(idx), fps=15)
+                    videoStreamer.cameraStatus.connect(handle_camera_status)
+                    print(f'{videoStreamer.id} connected')
+                    videoStreamers.append(videoStreamer)
+                core.set_video_sources(videoStreamers)
+                # videoStreamers = []
+                # for idx, rtsp_link in enumerate(rtsp_links):
+                #     videoStreamer = GstVideoStreamer(rtsp_link, [1920, 1080, 3], str(idx), fps=15)
+                #     videoStreamer.cameraStatus.connect(handle_camera_status)
+                #     print(f'{videoStreamer.id} connected')
+                #     videoStreamers.append(videoStreamer)
+                #     core.set_video_sources(videoStreamers)
+            else:
+                sources = findUsbCams()
+                if len(sources) >= 2 :
+                    print("set camera source to captutre card")
+                    core.set_video_sources(findUsbCams())
 
     import socket
 
@@ -156,24 +256,32 @@ if __name__ == '__main__':
                 print(f"Received discovery request from {addr}")
                 sock.sendto(b"SERVER_RESPONSE", addr)
                 break
+
+        sock.close()
         return addr
 
     cl_ip, _ = start_discovery_service(12345)
     print(cl_ip)
 
+    global manager
     manager = Manager(f"tcp://{cl_ip}:5558", f"tcp://{cl_ip}:5557")
     manager.start(manager_callback)
 
     def gotNewFrame(bboxes, id_, isDetection, ctime):
+        #print(f"Got new frame, bboxes : {bboxes} Id: {id} Is detection {isDetection}")
         m = Message()
         m.msgType = MessageType.MESSAGE_TYPE_IMAGE
         m.image.timestamp = int(ctime.value)
         for bbox in bboxes:
+            # Skip if bbox is None, doesn't have exactly 4 elements, or contains None values
+            if bbox is None or len(bbox) != 4 or not all(element is not None for element in bbox):
+                continue
+
+
+            # Add the bounding box to the image
             box = m.image.boxes.add()
-            box.x = bbox[0]
-            box.y = bbox[1]
-            box.w = bbox[2]
-            box.h = bbox[3]
+            box.x, box.y, box.w, box.h = bbox
+
         m.image.camType = id_
         if isDetection:
             m.image.trackMode = TrackMode.TRACK_MODE_DETECT
@@ -193,12 +301,13 @@ if __name__ == '__main__':
     core.newFrame.connect(gotNewFrame)
     core.coordsUpdated.connect(gotCoords)
 
-    # Set up the QTimer to check client connection every 10 seconds
-    timer = QTimer()
-    timer.timeout.connect(check_client_connection)
-    timer.start(10000)  # 10 seconds
+    # Start the connection thread
+    connection_thread = ConnectionThread()
+    connection_thread.start()
 
     try:
         app.exec_()
     except KeyboardInterrupt:
-        sys.exit(0)
+        connection_thread.stop()
+        connection_thread.join()
+        sys.exit(0)

config.yaml

@@ -6,4 +6,4 @@ rtsp_links: [
 thickness: 2
 
 
-debug: false
+debug: False

core.py

@@ -1,34 +1,40 @@
-#import os
-#os.environ['YOLO_VERBOSE'] = "false"
+import datetime
+import threading
 
+import pygame
+import cv2
+import numpy as np
+import torch
 from threading import Event, Thread
 from typing import List
-
-import numpy as np
 from PyQt5.QtCore import QThread, pyqtSlot, pyqtSignal, QUrl, QDir, pyqtProperty
-#from icecream import ic
-
+from matplotlib import pyplot as plt
 from detector import Detector
 from detector.utils import get_bbox_by_point
 from tracker import Tracker
 from video_streamer.videostreamer import VideoStreamer
 from time import sleep
 import time
-
 from PyQt5.QtCore import QObject, pyqtSignal
 import ctypes
 from ctypes import c_int64
+from server import rtsp_server, run_server, RTSPServer
+
 
+showTrack = False
 
 class Core(QThread):
-    newFrame = pyqtSignal(object, int, bool,ctypes.c_int64)
+    newFrame = pyqtSignal(object, int, bool, ctypes.c_int64)
     coordsUpdated = pyqtSignal(int, object, bool)
 
-    def __init__(self, video_sources: List[VideoStreamer], parent=None):
+    def __init__(self, video_sources, tracker=None, detector=None, parent=None):
         super(QThread, self).__init__(parent)
 
-        self.__detector = Detector(classes=[0, 2, 5, 7])
-        self.__tracker = Tracker()
+        self.__detector = detector
+        self.__tracker = tracker
+
+        self.__rtspserver = rtsp_server
+        threading.Thread(target=run_server, daemon=True).start()
 
         self.__video_sources = video_sources
         self.__processing_source = video_sources[0]
@@ -44,7 +50,12 @@ class Core(QThread):
         self.__tracking_thread = None
 
         self.__processing_id = 0
-        # ic()
+
+        self.__frame = None  # Frame property for Pygame
+
+    @pyqtProperty(np.ndarray)
+    def frame(self):
+        return self.__frame
 
     def set_thickness(self, thickness: int):
         self.__thickness = thickness
@@ -53,13 +64,15 @@ class Core(QThread):
         self.__processing_source = self.__video_sources[source_id]
         self.__processing_id = source_id
 
-    def set_video_sources(self, video_sources: List[VideoStreamer]):
-        self.__video_sources = video_sources
-        self.set_source(0)
+    def set_video_sources(self, video_sources):
+        if len(video_sources) >= 2:
+            self.__video_sources = video_sources
+            self.set_source(0)
 
     def __detection(self):
         while self.__is_detecting:
             try:
+                torch.cuda.empty_cache()
                 source = self.__processing_source
                 roi = self.__detection_roi
                 frame = source.get_frame()
@@ -71,10 +84,8 @@ class Core(QThread):
                     bbox = result[1:]
                     bbox[:2] += roi[:2]
                     global_bboxes.append(bbox)
-                    # color = (0, 0, 255) if cls == 0 else (80, 127, 255)
-                    # self.__draw_bbox(frame, bbox, color)
 
-                self.newFrame.emit(global_bboxes, self.__processing_id, True)
+                self.newFrame.emit(global_bboxes, self.__processing_id, True, c_int64(int(time.time() * 1e3)))
                 self.__detection_bboxes = np.array(global_bboxes)
                 self.__detection_frame = frame.copy()
                 sleep(0.03)
@@ -82,21 +93,96 @@ class Core(QThread):
                 print(e)
                 sleep(0.1)
 
+
+
     def __tracking(self):
         source = self.__processing_source
+        if showTrack:
+            pygame.init()
+
+            # Get actual screen resolution
+            info = pygame.display.Info()
+            screen_width, screen_height = info.current_w, info.current_h
+            screen = pygame.display.set_mode((screen_width, screen_height), pygame.FULLSCREEN)
+            pygame.display.set_caption('Tracking Frame')
+
+            clock = pygame.time.Clock()  # Add a clock to control frame rate
+
         while self.__is_tracking:
+            if showTrack:
+                for event in pygame.event.get():  # Prevent freezing by handling events
+                    if event.type == pygame.QUIT:
+                        pygame.quit()
+                        return
+
             ctime = c_int64(int(time.time() * 1000))  # Convert to c_int64
+
             frame = source.get_frame()
             bbox, success = self.__tracker.update(frame)
-            center = None
+
+
             if bbox is not None:
                 center = bbox[:2] + bbox[2:] // 2
                 self.coordsUpdated.emit(self.__processing_id, center, success)
                 self.newFrame.emit([bbox], self.__processing_id, False, ctime)
-                sleep(0.01)
-            else:
-                self.newFrame.emit([bbox], self.__processing_id, False, ctime)
-                sleep(0.05)
+
+                x, y, w, h = map(int, bbox)
+                box_color = (0, 255, 0) if success else (255, 0, 0)
+                cv2.rectangle(frame, (x, y), (x + w, y + h), box_color, 2)
+
+                self.__rtspserver.update_frame(frame)
+                if showTrack:
+
+                    # Convert OpenCV frame (BGR) to RGB
+                    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+                    font = cv2.FONT_HERSHEY_SIMPLEX
+                    font_scale = 2.25
+                    font_color = (255, 255, 0)
+                    thickness = 6
+                    position = (50, 450)  # Bottom-left corner of the text in the image
+                    now = datetime.datetime.now()
+                    time_string = now.strftime("%H:%M:%S.%f")[:-3]
+
+                    # 4. Use cv2.putText() to write time on the image
+                    cv2.putText(frame, time_string, position, font, font_scale, font_color, thickness, cv2.LINE_AA)
+                    cv2.putText(frame, f"{ctime}", (50, 380), font, font_scale, (255,255,255), thickness, cv2.LINE_AA)
+                    # print(ctime)
+
+                    frame = cv2.flip(frame, 1)  # Flip horizontally
+
+
+
+
+
+                    # Resize frame while maintaining aspect ratio
+                    frame_height, frame_width, _ = frame.shape
+                    aspect_ratio = frame_width / frame_height
+
+                    if aspect_ratio > (screen_width / screen_height):  # Wider than screen
+                        new_width = screen_width
+                        new_height = int(screen_width / aspect_ratio)
+                    else:  # Taller than screen
+                        new_height = screen_height
+                        new_width = int(screen_height * aspect_ratio)
+
+                    resized_frame = cv2.resize(frame, (new_width, new_height))
+
+                    # Convert to Pygame surface without unnecessary rotation
+                    frame_surface = pygame.surfarray.make_surface(resized_frame)
+
+                    # Optional: If rotation is needed, use pygame.transform.rotate()
+                    frame_surface = pygame.transform.rotate(frame_surface, -90)  # Example rotation
+
+                    # Center the frame
+                    x_offset = (screen_width - new_width) // 2
+                    y_offset = (screen_height - new_height) // 2
+
+                    screen.fill((0, 0, 0))  # Clear screen
+                    screen.blit(frame_surface, (x_offset, y_offset))
+                    pygame.display.flip()
+
+                    clock.tick(30)  # Limit FPS to prevent excessive CPU usage
 
     def start_detect(self, x: int, y: int, w: int, h: int):
         self.__detection_roi = [x, y, x + w, y + h]
@@ -110,51 +196,53 @@ class Core(QThread):
 
     def stop_detection(self):
         self.__is_detecting = False
-        if self.__detection_thread is not None:
+        if self.__detection_thread is not None and self.__detection_thread.is_alive():
             self.__detection_thread.join()
 
         self.__detection_thread = None
 
     def start_track(self, x: int, y: int, w: int = 0, h: int = 0):
-      try:
-        self.__is_detecting = False
-        self.__is_tracking = False
-        bbox = None
-        if w == 0:
-            if len(self.__detection_bboxes):
-                bbox = get_bbox_by_point(self.__detection_bboxes, np.array([x, y]))
-            frame = self.__detection_frame
-        else:
-            bbox = np.array([x, y, w, h])
-            frame = self.__processing_source.get_frame()
+        print(f"start tracking: {x}, {y}, {w}, {h}")
+        try:
+            self.__is_detecting = False
+            self.__is_tracking = False
+            bbox = None
+            if w == 0:
+                if len(self.__detection_bboxes):
+                    bbox = get_bbox_by_point(self.__detection_bboxes, np.array([x, y]))
+                frame = self.__detection_frame
+            else:
+                bbox = np.array([x, y, w, h])
+                frame = self.__processing_source.get_frame()
 
-        self.__tracker.stop()
+            self.__tracker.stop()
 
-        if bbox is not None:
-            self.__tracker.init(frame, bbox)
-        else:
+            if bbox is not None:
+                self.__tracker.init(frame, bbox)
+            else:
+                return
+        except Exception as e:
+            print(e)
             return
-      except Exception as e:
-        print(e)
-        return
 
-      if self.__tracking_thread is not None:
-          self.__tracking_thread.join()
-      self.__is_tracking = True
-      self.__tracking_thread = Thread(target=self.__tracking)
-      self.__tracking_thread.start()
-      sleep(0.03)
+        if self.__tracking_thread is not None:
+            self.__tracking_thread.join()
+            self.stop_track()
+        self.__is_tracking = True
+        self.__tracking_thread = Thread(target=self.__tracking)
+        self.__tracking_thread.start()
+        sleep(0.03)
 
     def stop_track(self):
+        if showTrack:
+            pygame.quit()
+        print("stop tracking")
         self.stop_detection()
         self.__tracker.stop()
         self.__is_tracking = False
-        if self.__tracking_thread is not None:
-            self.__tracking_thread.join()
-
         self.__tracking_thread = None
 
     def __draw_bbox(self, img: np.ndarray, bbox, color):
         thickness = self.__thickness
-        # cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2] + bbox[0], bbox[3] + bbox[1]),
-        #               color, thickness)
+        cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2] + bbox[0], bbox[3] + bbox[1]),
+                      color, thickness)

cvStreamer.py

@@ -0,0 +1,52 @@
+import cv2
+from pywt.version import release
+
+
+class cvStreamer():
+    def __init__(self, idx):
+        self.cap = cv2.VideoCapture(idx)
+        self.idx = idx
+
+    def isOpened(self):
+        isOpen = self.cap.isOpened()
+
+        if not isOpen:
+            self.release()
+        else:
+            print(f"usb cam open at {self.idx}")
+
+
+        return isOpen
+
+    def release(self):
+        self.cap.release()
+
+    def get_frame(self):
+        _, frame = self.cap.read()
+
+        # Get the original dimensions of the frame
+        height, width = frame.shape[:2]
+
+        # Define the maximum dimensions
+        max_width = 1920
+        max_height = 1080
+
+        # Calculate the aspect ratio
+        aspect_ratio = width / height
+
+        # Resize the frame if it exceeds the maximum dimensions
+        if width > max_width or height > max_height:
+            if aspect_ratio > 1:  # Landscape orientation
+                new_width = max_width
+                new_height = int(new_width / aspect_ratio)
+            else:  # Portrait orientation
+                new_height = max_height
+                new_width = int(new_height * aspect_ratio)
+
+            # Resize the frame
+            frame = cv2.resize(frame, (new_width, new_height))
+
+        return frame
+
+    def __del__(self):
+        self.release()

detector/demo.py

@@ -7,7 +7,7 @@ from utils import get_bbox_by_point
 if __name__ == '__main__':
     detector = Detector(classes=[0, 2, 5, 7])
 
-    cap = cv2.VideoCapture(1)
+    cap = cv2.VideoCapture(0)
     display_name = 'detector'
     cv2.namedWindow(display_name, cv2.WINDOW_NORMAL)
     cv2.setWindowProperty(display_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)

gpuMonitor.py

@@ -0,0 +1,47 @@
+import pynvml
+import time
+from colorama import Fore, Style, init
+import os
+
+# Initialize colorama
+init(autoreset=True)
+
+
+
+def monitor_gpu_ram_usage(interval=2, threshold_gb=2):
+    pynvml.nvmlInit()
+    # Initialize NVML
+    try:
+        device_count = pynvml.nvmlDeviceGetCount()
+        print(f"Found {device_count} GPU(s).")
+
+        while True:
+            for i in range(device_count):
+                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+                info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+
+
+
+                print(f"GPU {i}:")
+                print(f"  Total RAM: {info.total / 1024 ** 2:.2f} MB")
+                if(info.used / 1024 ** 2 >= 2.5 * 1024 ):
+                    print(Fore.RED + f"  Used RAM: {info.used / 1024 ** 2:.2f} MB")
+                    os.system("aplay /home/rog/repos/Tracker/NE-Smart-Tracker/Oxygen-Sys-Warning.wav")
+                else:
+                    print(f"  Used RAM: {info.used / 1024 ** 2:.2f} MB")
+                print(f"  Free RAM: {info.free / 1024 ** 2:.2f} MB")
+                print(Fore.GREEN + "-" * 30)
+                print(Fore.GREEN)
+
+            time.sleep(interval)  # Wait for the specified interval before checking again
+
+    except KeyboardInterrupt:
+        print("Monitoring stopped by user.")
+
+    finally:
+        # Shutdown NVML
+        pynvml.nvmlShutdown()
+
+if __name__ == "__main__":
+    monitor_gpu_ram_usage(interval=2, threshold_gb=2)  # Check every 2 seconds, threshold is 2 GB
+

message_queue/proto/SetCameraMessage.proto

@@ -3,4 +3,5 @@ syntax = "proto3";
 message SetCameraMessage {
   string ip = 1;
   int32 port = 2;
+  int32 cameraSource = 3;
 }

message_queue/proto/SetCameraMessage_pb2.py

@@ -19,7 +19,7 @@ DESCRIPTOR = _descriptor.FileDescriptor(
   syntax='proto3',
   serialized_options=None,
   create_key=_descriptor._internal_create_key,
-  serialized_pb=b'\n\x16SetCameraMessage.proto\",\n\x10SetCameraMessage\x12\n\n\x02ip\x18\x01 \x01(\t\x12\x0c\n\x04port\x18\x02 \x01(\x05\x62\x06proto3'
+  serialized_pb=b'\n\x16SetCameraMessage.proto\"B\n\x10SetCameraMessage\x12\n\n\x02ip\x18\x01 \x01(\t\x12\x0c\n\x04port\x18\x02 \x01(\x05\x12\x14\n\x0c\x63\x61meraSource\x18\x03 \x01(\x05\x62\x06proto3'
 )
 
 
@@ -47,6 +47,13 @@ _SETCAMERAMESSAGE = _descriptor.Descriptor(
       message_type=None, enum_type=None, containing_type=None,
       is_extension=False, extension_scope=None,
       serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='cameraSource', full_name='SetCameraMessage.cameraSource', index=2,
+      number=3, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
   ],
   extensions=[
   ],
@@ -60,7 +67,7 @@ _SETCAMERAMESSAGE = _descriptor.Descriptor(
   oneofs=[
   ],
   serialized_start=26,
-  serialized_end=70,
+  serialized_end=92,
 )
 
 DESCRIPTOR.message_types_by_name['SetCameraMessage'] = _SETCAMERAMESSAGE

server.py

@@ -0,0 +1,112 @@
+import cv2
+import gi
+import threading
+import numpy as np
+
+gi.require_version('Gst', '1.0')
+gi.require_version('GstRtspServer', '1.0')
+from gi.repository import Gst, GstRtspServer, GLib
+import socket
+
+def get_local_ip():
+    try:
+        # Create a socket object
+        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+
+        # Connect to an external server (doesn't actually send any data)
+        s.connect(("8.8.8.8", 80))
+
+        # Get the local IP address
+        local_ip = s.getsockname()[0]
+
+        # Close the socket
+        s.close()
+
+        return local_ip
+    except Exception as e:
+        return str(e)
+
+
+
+class VideoStream(GstRtspServer.RTSPMediaFactory):
+    def __init__(self, fps=25, width=1920, height=1080):
+        super(VideoStream, self).__init__()
+        self.fps = fps
+        self.width = width
+        self.height = height
+        self.frame = np.zeros((height, width, 3), dtype=np.uint8)  # Black frame initially
+        self.lock = threading.Lock()
+
+    def update_frame(self, frame):
+        """Externally updates the current frame."""
+        with self.lock:
+            self.frame = cv2.resize(frame, (self.width, self.height))  # Resize if necessary
+
+    def on_need_data(self, src, length):
+        """Provides frames when the pipeline requests data."""
+        with self.lock:
+            frame_rgb = cv2.cvtColor(self.frame, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
+            data = frame_rgb.tobytes()  # Convert to bytes
+
+        buf = Gst.Buffer.new_allocate(None, len(data), None)
+        buf.fill(0, data)
+        buf.duration = Gst.SECOND // self.fps
+        buf.pts = buf.dts = Gst.CLOCK_TIME_NONE
+        src.emit("push-buffer", buf)
+
+    def do_create_element(self, url):
+        """Creates the GStreamer pipeline for RTSP streaming."""
+        pipeline_str = (
+            "appsrc name=source is-live=true format=GST_FORMAT_TIME "
+            "caps=video/x-raw,format=RGB,width={},height={},framerate={}/1 "
+            "! videoconvert ! video/x-raw,format=I420 "
+            "! nvh264enc preset=low-latency rc-mode=cbr bitrate=2500 gop-size=7 zerolatency=true "
+            "! h264parse ! rtph264pay config-interval=1 name=pay0 pt=96"
+        ).format(self.width, self.height, self.fps)
+
+
+        pipeline = Gst.parse_launch(pipeline_str)
+        src = pipeline.get_by_name("source")
+        src.connect("need-data", self.on_need_data)
+        return pipeline
+
+
+class RTSPServer:
+    def __init__(self, ip="0.0.0.0", port=8554, mount_point="/stream"):
+        Gst.init(None)
+        self.server = GstRtspServer.RTSPServer()
+        self.server.set_address(ip)
+        self.server.set_service(str(port))
+
+        self.video_stream = VideoStream()
+        self.server.get_mount_points().add_factory(mount_point, self.video_stream)
+        self.server.attach(None)
+
+    def start(self):
+        print(f"RTSP Server running at rtsp://{self.server.get_address()}:{self.server.get_service()}/stream")
+        loop = GLib.MainLoop()
+        loop.run()
+
+    def update_frame(self, frame):
+        """Externally updates the current frame for streaming."""
+        self.video_stream.update_frame(frame)
+
+# Global server instance
+rtsp_server = RTSPServer(get_local_ip(),41231)
+
+def run_server():
+    rtsp_server.start()
+
+# def stream_webcam():
+#     cap = cv2.VideoCapture("/home/mht/Downloads/bcd2890d71caaf0e095b95c9b525973f61186656-360p.mp4")  # Open webcam
+#     while cap.isOpened():
+#         ret, frame = cap.read()
+#         if ret:
+#             rtsp_server.update_frame(frame)  # Send frame to RTSP server
+
+# if __name__ == "__main__":
+#     # Start RTSP server in a separate thread
+#     threading.Thread(target=run_server, daemon=True).start()
+#
+#     # Stream webcam frames
+#     stream_webcam()

tracker/ltr/external/PreciseRoIPooling/pytorch/prroi_pool/functional.py

@@ -12,6 +12,7 @@
 import torch
 import torch.autograd as ag
 
+
 __all__ = ['prroi_pool2d']
 
 

tracker/ltr/models/backbone/resnet.py

@@ -6,6 +6,7 @@ from torchvision.models.resnet import model_urls
 from .base import Backbone
 
 
+
 class Bottleneck(nn.Module):
     expansion = 4
 
@@ -22,6 +23,7 @@ class Bottleneck(nn.Module):
         self.downsample = downsample
         self.stride = stride
 
+
     def forward(self, x):
         residual = x
 

tracker/ltr/models/bbreg/atom_iou_net.py

@@ -2,6 +2,7 @@ import torch.nn as nn
 import torch
 from ltr.models.layers.blocks import LinearBlock
 from ltr.external.PreciseRoIPooling.pytorch.prroi_pool import PrRoIPool2D
+torch.cuda.empty_cache()
 
 
 def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):

tracker/ltr/models/layers/distance.py

@@ -14,12 +14,14 @@ class DistanceMap(nn.Module):
         super().__init__()
         self.num_bins = num_bins
         self.bin_displacement = bin_displacement
+        torch.cuda.empty_cache()
 
     def forward(self, center, output_sz):
         """Create the distance map.
         args:
             center: Torch tensor with (y,x) center position. Dims (batch, 2)
             output_sz: Size of output distance map. 2-dimensional tuple."""
+        torch.cuda.empty_cache()
 
         center = center.view(-1,2)
 

tracker/ltr/models/target_classifier/features.py

@@ -4,9 +4,11 @@ from ltr.models.layers.normalization import InstanceL2Norm
 
 
 
+
 def residual_bottleneck(feature_dim=256, num_blocks=1, l2norm=True, final_conv=False, norm_scale=1.0, out_dim=None,
                         interp_cat=False, final_relu=False, final_pool=False, input_dim=None, final_stride=1):
     """Construct a network block based on the Bottleneck block used in ResNet."""
+
     if out_dim is None:
         out_dim = feature_dim
     if input_dim is None:

tracker/pytracking/features/augmentation.py

@@ -6,16 +6,18 @@ import cv2 as cv
 import random
 from pytracking.features.preprocessing import numpy_to_torch, torch_to_numpy
 
-
 class Transform:
     """Base data augmentation transform class."""
 
     def __init__(self, output_sz = None, shift = None):
         self.output_sz = output_sz
         self.shift = (0,0) if shift is None else shift
+        torch.cuda.empty_cache()
 
 
     def crop_to_output(self, image):
+        torch.cuda.empty_cache()
+
         if isinstance(image, torch.Tensor):
             imsz = image.shape[2:]
             if self.output_sz is None:
@@ -67,6 +69,7 @@ class Rotate(Transform):
         super().__init__(output_sz, shift)
         self.angle = math.pi * angle/180
 
+
     def __call__(self, image, is_mask=False):
         if isinstance(image, torch.Tensor):
             return self.crop_to_output(numpy_to_torch(self(torch_to_numpy(image))))
@@ -90,6 +93,7 @@ class Blur(Transform):
         self.filter[0] = self.filter[0].view(1,1,-1,1) / self.filter[0].sum()
         self.filter[1] = self.filter[1].view(1,1,1,-1) / self.filter[1].sum()
 
+
     def __call__(self, image, is_mask=False):
         if isinstance(image, torch.Tensor):
             sz = image.shape[2:]

tracker/pytracking/features/preprocessing.py

@@ -4,10 +4,12 @@ import numpy as np
 
 
 def numpy_to_torch(a: np.ndarray):
+    torch.cuda.empty_cache()
     return torch.from_numpy(a).float().permute(2, 0, 1).unsqueeze(0)
 
 
 def torch_to_numpy(a: torch.Tensor):
+    torch.cuda.empty_cache()
     return a.squeeze(0).permute(1,2,0).numpy()
 
 
@@ -20,7 +22,7 @@ def sample_patch_transformed(im, pos, scale, image_sz, transforms, is_mask=False
         image_sz: Size to resize the image samples to before extraction.
         transforms: A set of image transforms to apply.
     """
-
+    torch.cuda.empty_cache()
     # Get image patche
     im_patch, _ = sample_patch(im, pos, scale*image_sz, image_sz, is_mask=is_mask)
 
@@ -39,6 +41,7 @@ def sample_patch_multiscale(im, pos, scales, image_sz, mode: str='replicate', ma
         mode: how to treat image borders: 'replicate' (default), 'inside' or 'inside_major'
         max_scale_change: maximum allowed scale change when using 'inside' and 'inside_major' mode
     """
+    torch.cuda.empty_cache()
     if isinstance(scales, (int, float)):
         scales = [scales]
 
@@ -62,7 +65,7 @@ def sample_patch(im: torch.Tensor, pos: torch.Tensor, sample_sz: torch.Tensor, o
         mode: how to treat image borders: 'replicate' (default), 'inside' or 'inside_major'
         max_scale_change: maximum allowed scale change when using 'inside' and 'inside_major' mode
     """
-
+    torch.cuda.empty_cache()
     # if mode not in ['replicate', 'inside']:
     #     raise ValueError('Unknown border mode \'{}\'.'.format(mode))
 

tracker/pytracking/libs/dcf.py

@@ -3,6 +3,7 @@ import torch
 
 def max2d(a: torch.Tensor) -> (torch.Tensor, torch.Tensor):
     """Computes maximum and argmax in the last two dimensions."""
+    torch.cuda.empty_cache()
 
     max_val_row, argmax_row = torch.max(a, dim=-2)
     max_val, argmax_col = torch.max(max_val_row, dim=-1)

tracker/pytracking/tracker/dimp/dimp.py

@@ -65,6 +65,7 @@ class DiMP():
 
 
     def track(self, image) -> dict:
+        torch.cuda.empty_cache()
 
         # Convert image
         im = numpy_to_torch(image)
@@ -213,7 +214,7 @@ class DiMP():
 
         # Compute augmentation size
         aug_expansion_factor = self.params.get('augmentation_expansion_factor', None)
-        ic(self.params.get('augmentation_expansion_factor', None))
+
 
 
         aug_expansion_sz = (self.img_sample_sz * aug_expansion_factor).long()

tracker/pytracking/utils/params.py

@@ -5,7 +5,7 @@ class TrackerParams:
     """Class for tracker parameters."""
 
     image_sample_size = 18 * 16
-    search_area_scale = 5
+    search_area_scale = 7
 
     # Learning parameters
     sample_memory_size = 50