Coverage for skcvideo/controlled_fps_video_capture.py: 75%

110 statements  

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

1import logging 

2import math 

3 

4import cv2 

5from ffmphisdp.pyav_reader import ControlledVideoReader as ControlledFPSVideoCapture_ffmphisdp 

6 

7 

8class ControlledFPSVideoCapture: 

9 FRAME_SELECTION_LEGACY = "legacy" 

10 FRAME_SELECTION_FFMPHISDP = "ffmphisdp" 

11 

12 def __new__(cls, *args, **kwargs): 

13 kwargs_copy = kwargs.copy() 

14 frame_selection_version = kwargs_copy.pop( 

15 "frame_selection_version", ControlledFPSVideoCapture.FRAME_SELECTION_LEGACY 

16 ) 

17 if frame_selection_version == ControlledFPSVideoCapture.FRAME_SELECTION_LEGACY: 

18 reader = ControlledFPSVideoCapture_skcvideo(*args, **kwargs_copy) 

19 reader.frame_selection_version = frame_selection_version 

20 return reader 

21 elif frame_selection_version == ControlledFPSVideoCapture.FRAME_SELECTION_FFMPHISDP: 

22 

23 def set_convert(self, prop_id, value): 

24 """Convert form global frame index to selection frame index""" 

25 if prop_id != cv2.CAP_PROP_POS_FRAMES: 

26 raise NotImplementedError 

27 idx_global = self.frame_selection_index_mapping[value] 

28 print(f"Automatically converting index from {value} to {idx_global} (local to global)") 

29 self.set_frame_idx(idx_global) 

30 

31 def get_convert(self, prop_id): 

32 """convert from global frame index to selection frame index""" 

33 if prop_id == cv2.CAP_PROP_POS_FRAMES: 

34 print(self.current_frame_index) 

35 return self.frame_selection_index_mapping.index(self.current_frame_index) 

36 elif prop_id == cv2.CAP_PROP_FRAME_COUNT: 

37 return len(self.frame_selection_index_mapping) 

38 else: 

39 raise NotImplementedError 

40 

41 ControlledFPSVideoCapture_ffmphisdp.set = set_convert 

42 ControlledFPSVideoCapture_ffmphisdp.get = get_convert 

43 reader = ControlledFPSVideoCapture_ffmphisdp(*args, **kwargs_copy) 

44 reader.frame_selection_version = frame_selection_version 

45 reader.frame_selection_index_mapping = [ 

46 idx_global for idx_global, selected in enumerate(reader.frame_selection) if selected 

47 ] 

48 return reader 

49 else: 

50 raise ValueError("Invalid frame_selection_version") 

51 

52 

53class ControlledFPSVideoCapture_skcvideo: 

54 """ 

55 VideoCapture Reader from any input fps to any output fps. 

56 Has been tested with output_fps=10 and input_fps=10, 25, 30, 50, 60 

57 Note that input frame and output frame are the last decoded frames 

58 While cap.get(cv2.CAP_PROP_POS_FRAMES) is next frame to be read 

59 

60 This class aims to simulate the result of downsampling by ffmpeg 

61 directly by reading the original video. This avoid having to store 

62 multiple versions of the video. 

63 """ 

64 

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

66 logger = kwargs.pop("logger", logging.getLogger(__name__)) 

67 

68 self.input_frame = -1 

69 self.output_frame = -1 

70 self.output_fps = kwargs.pop("fps", 10) 

71 

72 self._cap = cv2.VideoCapture(*args, **kwargs) 

73 # We round fps cause sometimes opencv gives fps like 25.00001429630351 while ffmpeg gives 25 

74 # Still don't know if we have correct behaviour for weird fps like 29.97 

75 self.input_fps = round(self._cap.get(cv2.CAP_PROP_FPS)) 

76 

77 if self.output_fps == "same": 

78 self.output_fps = self.input_fps 

79 

80 message = ( 

81 f"Initializing ControlledFPSVideoCapture with input_fps={self.input_fps} and output_fps={self.output_fps}" 

82 ) 

83 logger.info(message) 

84 

85 def read(self): 

86 self.output_frame += 1 

87 if self.output_fps == self.input_fps: 

88 self.input_frame += 1 

89 return self._cap.read() 

90 

91 # This general case do not work with output_fps == input_fps (weird case output_frame == 2) 

92 if self.output_frame < 2: 

93 # Do not understand why but that's what ffmpeg does 

94 frames_to_get = 1 

95 else: 

96 frames_to_get = int( 

97 math.ceil((self.output_frame + 1) * float(self.input_fps) / self.output_fps) 

98 - math.ceil((self.output_frame) * float(self.input_fps) / self.output_fps) 

99 ) 

100 # Do not understand why but that's what ffmpeg does 

101 if self.output_frame == 2: 

102 frames_to_get -= 2 

103 

104 for _ in range(frames_to_get): 

105 self.input_frame += 1 

106 ret, im = self._cap.read() 

107 if ret is False: 

108 return ret, im 

109 

110 return ret, im 

111 

112 def get(self, prop_id): 

113 if prop_id == cv2.CAP_PROP_BUFFERSIZE: 

114 raise NotImplementedError 

115 elif prop_id == cv2.CAP_PROP_FPS: 

116 return self.output_fps 

117 elif prop_id == cv2.CAP_PROP_POS_FRAMES: 

118 return self.output_frame + 1 

119 elif prop_id == cv2.CAP_PROP_POS_MSEC: 

120 return 1000 * float(self.output_frame + 1) / self.output_fps 

121 elif prop_id == cv2.CAP_PROP_FRAME_COUNT: 

122 return int(math.floor(float(self._cap.get(prop_id)) / self.input_fps * self.output_fps)) 

123 else: 

124 return self._cap.get(prop_id) 

125 

126 def set(self, prop_id, value): 

127 if prop_id in [cv2.CAP_PROP_BUFFERSIZE, cv2.CAP_PROP_FPS, cv2.CAP_PROP_FRAME_COUNT]: 

128 raise NotImplementedError 

129 elif prop_id in [cv2.CAP_PROP_POS_FRAMES, cv2.CAP_PROP_POS_MSEC, cv2.CAP_PROP_POS_AVI_RATIO]: 

130 if prop_id == cv2.CAP_PROP_POS_FRAMES: 

131 if value < 0: 

132 value = 0 

133 output_frame = value - 1 

134 elif prop_id == cv2.CAP_PROP_POS_MSEC: 

135 output_frame = round(float(value * self.output_fps) / 1000) - 1 

136 elif prop_id == cv2.CAP_PROP_POS_AVI_RATIO: 

137 output_frame = round(self.get(cv2.CAP_PROP_FRAME_COUNT) * cv2.CAP_PROP_POS_AVI_RATIO) - 1 

138 self.output_frame = output_frame 

139 fps_ratio = float(self.input_fps) / self.output_fps 

140 if fps_ratio == 1.0 or output_frame < 2: 

141 self.input_frame = self.output_frame 

142 else: 

143 self.input_frame = math.floor(output_frame * fps_ratio) - math.floor(fps_ratio + 1) 

144 return self._cap.set(prop_id, self.input_frame + 1) 

145 else: 

146 return self._cap.set(prop_id, value) 

147 

148 def generator(self): 

149 while True: 

150 ret, im = self.read() 

151 if ret is False: 

152 raise StopIteration 

153 yield im 

154 

155 def release(self): 

156 return self._cap.release() 

157 

158 def isOpened(self): 

159 return self._cap.isOpened() 

160 

161 def retrieve(self): 

162 return self._cap.retrieve() 

163 

164 def open(self, *args, **kwargs): 

165 """Could be implemented but we don't use it""" 

166 raise NotImplementedError 

167 

168 def grab(self): 

169 """Could be implemented but we don't use it""" 

170 raise NotImplementedError