Coverage for skcvideo/core.py: 0%

124 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-10-02 14:10 +0200

1import os 

2import sys 

3import time 

4 

5import cv2 

6import imageio 

7import numpy as np 

8 

9from skcvideo.colors import BLACK 

10from skcvideo.utils import put_text 

11 

12 

13class Button: 

14 """ 

15 Used to define a clickable button on the image executing a given callback 

16 when cliked. Some data specifying the button can be passed at the object 

17 creation. 

18 

19 Args: 

20 - hitbox: tuple (x1, y1, x2, y2) the bounding box of the clickable area. 

21 - callback: a function taking x, y (the coordinates of the click) and 

22 optionnaly data as arguments. 

23 - data (optionnal): data of any shape that will be used by the callback. 

24 """ 

25 

26 def __init__(self, hitbox, callback, data=None): 

27 self.hitbox = hitbox 

28 self.data = data 

29 self.given_callback = callback 

30 

31 def callback(self, *kwargs): 

32 if self.data is None: 

33 return self.given_callback(*kwargs) 

34 else: 

35 return self.given_callback(self.data, *kwargs) 

36 

37 

38class Reader: 

39 """ 

40 A video displayer that allows interaction with the image by using buttons 

41 or keyboard. 

42 

43 The main advantage of this displayer is that it allows to read the video 

44 backward while keeping relatively fast. 

45 

46 The best way to use this displayer is to make your own class inheriting 

47 from this one and overridding its methods. 

48 """ 

49 

50 def __init__(self, *args, **kwargs): 

51 self.to_exit = False 

52 self.size = kwargs.get("size", (1920, 1080)) 

53 self.min_frame = kwargs.get("min_frame", 0) 

54 self.max_frame = kwargs.get("max_frame", 9000) 

55 self.frame = self.min_frame 

56 

57 self.is_playing = False 

58 self.max_playing_fps = kwargs.get("max_playing_fps", 10.0) 

59 self.playing_fps = None 

60 

61 # The key/function mapping 

62 self.keydict = { 

63 "k": self.next, 

64 "j": self.previous, 

65 "q": self.exit, 

66 "p": self.toggle_is_playing, 

67 " ": self.toggle_is_playing, 

68 } 

69 

70 # Widgets (the order of the widgets defines the order in which they 

71 # will be drawn) 

72 self.widgets = [] 

73 for widget in kwargs.get("widgets", []): 

74 self.add_widget(widget) 

75 

76 # The clickable buttons 

77 self.buttons = [] 

78 

79 self.background = self.build() 

80 

81 self._refresh() 

82 

83 @property 

84 def image_to_disp(self): 

85 """ 

86 This property specifies the image to be displayed. You would override 

87 it at your convenience e.g. to only display a subpart of the global 

88 image. 

89 """ 

90 return self.big_image 

91 

92 def toggle_is_playing(self): 

93 self.is_playing = not self.is_playing 

94 

95 def next(self): 

96 if self.frame < self.max_frame: 

97 self.frame += 1 

98 self._refresh() 

99 

100 def previous(self): 

101 if self.frame > self.min_frame: 

102 self.frame -= 1 

103 self._refresh() 

104 

105 def jump(self, frame): 

106 if frame < self.min_frame: 

107 self.frame = self.min_frame 

108 elif frame >= self.max_frame: 

109 self.frame = self.max_frame - 1 

110 else: 

111 self.frame = frame 

112 

113 def exit(self): 

114 self.to_exit = True 

115 

116 def add_widget(self, widget): 

117 self.widgets.append(widget) 

118 widget.parent = self 

119 

120 def build(self): 

121 """ 

122 Here you define the elements of the image that don't change throughout 

123 the video or manipulations. 

124 """ 

125 im = np.zeros((self.size[1], self.size[0], 3), dtype=np.uint8) 

126 for widget in self.widgets: 

127 widget.build(im) 

128 return im 

129 

130 def click_event(self, event, x, y, flags, param): 

131 """ 

132 Part of the core engine that manages the buttons. 

133 

134 /!\\ Should not be overridden without knowing what you do. 

135 """ 

136 if event == cv2.EVENT_LBUTTONUP: 

137 for button in [b for widget in [self] + self.widgets for b in getattr(widget, "buttons", [])]: 

138 x1, y1, x2, y2 = button.hitbox 

139 if x1 <= x < x2 and y1 <= y < y2: 

140 button.callback(x, y) 

141 self._refresh() 

142 

143 def _refresh(self): 

144 """ 

145 Here you define the appearance of the image to be displayed with 

146 respect to structural elements such as the frame index. 

147 

148 It is called each time the user is interacting with the image 

149 (clicks, keys, previous, next, ...) to allow updating it with new 

150 information. 

151 """ 

152 self.big_image = self.background.copy() 

153 self.refresh() 

154 

155 def refresh(self): 

156 for widget in self.widgets: 

157 widget.refresh(self.big_image, self.frame) 

158 

159 put_text( 

160 img=self.big_image, 

161 text=f"Frame {self.frame}", 

162 org=(20, 20), 

163 align_x="left", 

164 align_y="top", 

165 color=BLACK, 

166 thickness=3, 

167 ) 

168 put_text( 

169 img=self.big_image, 

170 text=f"Frame {self.frame}", 

171 org=(20, 20), 

172 align_x="left", 

173 align_y="top", 

174 ) 

175 if self.is_playing and self.playing_fps is not None: 

176 put_text( 

177 img=self.big_image, 

178 text=f"fps: {self.playing_fps:.2f}", 

179 org=(1900, 20), 

180 align_x="right", 

181 align_y="top", 

182 ) 

183 

184 def start(self): 

185 """ 

186 Part of the core engine that manages the display of the image and the 

187 keys. 

188 

189 /!\\ Should not be overridden without knowing what you do. 

190 """ 

191 cv2.namedWindow("image", cv2.WINDOW_NORMAL) 

192 cv2.resizeWindow("image", 1280, 720) 

193 cv2.setMouseCallback("image", self.click_event) 

194 

195 last_time = None 

196 while not self.to_exit: 

197 cv2.imshow("image", self.image_to_disp) 

198 key = cv2.waitKey(1) & 0xFF 

199 

200 for k, fun in self.keydict.items(): 

201 if key == ord(k): 

202 fun() 

203 

204 if self.is_playing: 

205 if last_time is not None: 

206 spent_time = time.time() - last_time 

207 if spent_time < 1 / self.max_playing_fps: 

208 time.sleep(1 / self.max_playing_fps - spent_time) 

209 

210 self.playing_fps = 1 / (time.time() - last_time) 

211 last_time = time.time() 

212 self.next() 

213 

214 def create_video(self, video_path="video.mp4", min_frame=None, max_frame=None, force_overwrite=False): 

215 if not force_overwrite and os.path.exists(video_path): 

216 print("video_path already exists, overwite (y/n)?") 

217 answer = input() 

218 if answer.lower() != "y": 

219 return 

220 video = imageio.get_writer(video_path, "ffmpeg", fps=10, quality=5.5) 

221 print("Creating video...") 

222 if min_frame is None: 

223 min_frame = self.min_frame 

224 if max_frame is None: 

225 max_frame = self.max_frame 

226 for frame in range(min_frame, max_frame): 

227 sys.stdout.write(f"\r{frame - min_frame}/{max_frame - min_frame - 1}") 

228 sys.stdout.flush() 

229 self.frame = frame 

230 self._refresh() 

231 video.append_data(cv2.cvtColor(self.big_image, cv2.COLOR_BGR2RGB)) 

232 sys.stdout.write("\n") 

233 sys.stdout.flush() 

234 print("Done") 

235 video.close() 

236 

237 

238if __name__ == "__main__": 

239 reader = Reader() 

240 reader.start()