version: v0.0.1

utils.py

@@ -0,0 +1,89 @@
+import numpy as np
+import time
+import cv2
+import demo_const as AI_CONST
+import copy
+
+from load_models import model_manager
+from ultralytics.data.loaders import LoadStreams,LOGGER
+
+def class_info():
+    od_model = model_manager.get_od()  # pt파일에 따라 값이 달라짐
+
+    class_dict = copy.deepcopy(od_model.model.names)
+    class_dict[1000] = 'signalman'
+    return class_dict
+
+
+CLASS_INFORMATION = {}
+if not CLASS_INFORMATION:
+    CLASS_INFORMATION = class_info()
+    CLASS_SWAP_INFO = {v:k for k,v in CLASS_INFORMATION.items()}
+    
+class LoadStreamsDaool(LoadStreams):
+    
+    def update(self, i, cap, stream):
+        """Read stream `i` frames in daemon thread."""
+        n, f = 0, self.frames[i]  # frame number, frame array
+        while self.running and cap.isOpened() and n < (f - 1):
+            if len(self.imgs[i]) < AI_CONST.LOADSTREAMS_IMG_BUFFER:  # keep a <=30-image buffer
+                n += 1
+                cap.grab()  # .read() = .grab() followed by .retrieve()
+                if n % self.vid_stride == 0:
+                    success, im = cap.retrieve()
+                    if not success:
+                        im = np.zeros(self.shape[i], dtype=np.uint8)
+                        LOGGER.warning("WARNING ⚠️ Video stream unresponsive, please check your IP camera connection.")
+                        cap.open(stream)  # re-open stream if signal was lost
+                    if self.buffer:
+                        self.imgs[i].append(im)
+                    else:
+                        self.imgs[i] = [im]
+            else:
+                time.sleep(0.01)  # wait until the buffer is empty
+
+class CustomVideoCapture:
+    def __init__(self) -> None:
+        self.fps = 30
+        # self.fourcc = cv2.VideoWriter_fourcc(*'X264')
+        self.fourcc = cv2.VideoWriter_fourcc(*'XVID')
+        self.size = None
+        self.video_writer = None
+
+    def set_frame_size(self, image):
+        if isinstance(image, np.ndarray):
+            height,width,ch = image.shape
+            self.size = (int(width), int (height))
+
+    def set_video_writer(self, path):
+        if self.size == None:
+            raise
+        
+        self.video_writer = cv2.VideoWriter(path, self.fourcc, self.fps, self.size)
+    
+    def write_video(self,frame):
+        if self.video_writer:
+            self.video_writer.write(frame)
+            
+    def release_video(self):
+        if self.video_writer:
+            self.video_writer.release()
+            self.video_writer = None
+
+
+def get_monitorsize():
+    width = 0
+    height = 0
+    try:
+        import pyautogui #pip install pyautogui
+        width = pyautogui.size()[0] # getting the width of the screen
+        height = pyautogui.size()[1] # getting the height of the screen
+    
+    except Exception:
+        (width,height) = AI_CONST.FHD_RESOLUTION
+        
+    finally:
+        return (width,height)
+    
+def img_resize(image,size):
+    return cv2.resize(image, dsize=size, interpolation=cv2.INTER_AREA)