89 lines
3.0 KiB
Python
89 lines
3.0 KiB
Python
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) |