Spaces:
Running
Running
| import os | |
| import cv2 | |
| import copy | |
| import time | |
| import numpy as np | |
| # Demo libs | |
| import spiga.demo.visualize.plotter as plt | |
| class Viewer: | |
| def __init__(self, window_title, width=None, height=None, fps=30): | |
| """ | |
| Initialization of the viewer canvas using width and height in pixels | |
| :param window_title: The string with the window title to display. | |
| :param width: The given width in pixels of the window canvas. | |
| :param height: The given height in pixels of the window canvas. | |
| :param fps: Frames per second | |
| """ | |
| # Visualizer parameters | |
| self.canvas = None | |
| self.width = width | |
| self.height = height | |
| self.window_title = window_title | |
| self.visualize = False | |
| # Time variables | |
| self.fps = fps | |
| self.fps_inference = 0 | |
| self.fps_mean = 0 | |
| self.fps_lifo = np.zeros(self.fps) | |
| self.timer = time.time() | |
| self.frame_cnt = -1 | |
| # Video/Image writer | |
| self.write = False | |
| self.video_name = window_title # Initial name | |
| self.video_path = None | |
| self.video_writer = None | |
| # Plots | |
| self.plotter = plt.Plotter() | |
| self.fps_draw_params = {'text_size': 0.75, | |
| 'text_thick': 2, | |
| 'coord': (10, 50), | |
| 'font': cv2.FONT_HERSHEY_SIMPLEX, | |
| 'color': (255, 255, 255)} | |
| def start_view(self): | |
| self._kill_window() | |
| cv2.namedWindow(self.window_title) | |
| self.visualize = True | |
| def record_video(self, video_path, video_name=None): | |
| self.write = True | |
| if video_name is not None: | |
| self.video_name = video_name | |
| self.video_path = video_path | |
| if not os.path.exists(video_path): | |
| os.makedirs(video_path) | |
| file_name = os.path.join(self.video_path, self.video_name + '.mp4') | |
| self.video_writer = cv2.VideoWriter(file_name, cv2.VideoWriter_fourcc(*'MP4V'), | |
| self.fps, (int(self.width), int(self.height))) | |
| def save_canvas(self, file_path=None): | |
| if file_path is None: | |
| if self.video_path is None: | |
| raise ValueError('Path not defined neither video_path is available') | |
| else: | |
| file_path = self.video_path | |
| file_name = os.path.join(file_path, '/%s_%i.jpg' % (self.video_name, self.frame_cnt)) | |
| cv2.imwrite(file_path + file_name, self.canvas) | |
| def reset_params(self, width, height, window_title, fps=30): | |
| self.width = width | |
| self.height = height | |
| self._kill_window() | |
| if self.video_name == self.window_title: | |
| self.video_name = window_title | |
| self.window_title = window_title | |
| self.fps = fps | |
| def close(self): | |
| if self.write: | |
| self.video_writer.release() | |
| self._kill_window() | |
| def process_image(self, input_img, drawers=(), show_attributes=('fps')): | |
| # Variables | |
| image = copy.copy(input_img) | |
| img_h, img_w, img_ch = image.shape | |
| # Convert gray scale image to color if needed | |
| if img_ch == 1: | |
| image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR) | |
| # Draw features on image | |
| image = self._draw_over_canvas(image, drawers, show_attributes) | |
| # Resize image if needed to canvas shape | |
| if img_w != self.width or img_h != self.height: | |
| image = cv2.resize(image, (self.width, self.height)) | |
| # Update canvas | |
| self.canvas = image | |
| # Visualize FPS | |
| if 'fps' in show_attributes: | |
| self._plot_fps() | |
| # Write the resulting frame | |
| if self.write: | |
| self.video_writer.write(self.canvas) | |
| # Timing loop variables | |
| loop_time = self._update_timers() | |
| break_flag = False | |
| # Visualization | |
| if self.visualize: | |
| cv2.imshow(self.window_title, self.canvas) | |
| sleep_time = int(1000 * (1 / self.fps - loop_time)) | |
| if sleep_time <= 0: | |
| sleep_time = 1 | |
| if cv2.waitKey(sleep_time) & 0xFF == ord('q'): | |
| break_flag = True | |
| self.timer = time.time() | |
| return break_flag | |
| def _plot_fps(self): | |
| # Plot algorithm time | |
| params = self.fps_draw_params | |
| cv2.putText(self.canvas, ('FPS: %.2f' % self.fps_mean), params['coord'], params['font'], params['text_size'], | |
| params['color'], params['text_thick'], cv2.LINE_AA) | |
| def _draw_over_canvas(self, image, drawers, show_attributes): | |
| for drawer in drawers: | |
| image = drawer.plot_features(image, self.plotter, show_attributes) | |
| return image | |
| def _kill_window(self): | |
| self.visualize = False | |
| try: | |
| cv2.destroyWindow(self.window_title) | |
| except: | |
| pass | |
| def _update_timers(self): | |
| self.frame_cnt += 1 | |
| loop_time = time.time() - self.timer | |
| self.fps_inference = 1/loop_time | |
| lifo_idx = self.frame_cnt % self.fps | |
| self.fps_lifo[lifo_idx] = self.fps_inference | |
| if lifo_idx == 0: | |
| self.fps_mean = np.mean(self.fps_lifo) | |
| return loop_time | |