Coverage for skcvideo/video_capture.py: 72%

78 statements  

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

1import cv2 

2 

3from skcvideo.controlled_fps_video_capture import ControlledFPSVideoCapture 

4 

5DEFAULT_WIDTH = 1280 

6DEFAULT_HEIGHT = 720 

7 

8SET_READ_TRADEOFF = 20 

9MAX_STORED = 100 

10AMONG_FRAMES = 20 

11 

12 

13class TransformedImageVideoCapture: 

14 """ 

15 This VideoCapture overlay aims to add operations on the images before 

16 providing them. 

17 

18 For instance, the operations are: 

19 - change channel order to switch from BGR to RGB 

20 - resize to desired shape 

21 

22 Additionnaly, you can give an initial frame for the VideoCapture 

23 

24 Args: 

25 - min_frame: the initial frame index of the video to start with 

26 - colormap: the order of the image channels (can be 'rgb' or 'bgr') 

27 - resize: tuple (width, height) or None: the shape of the outputted 

28 image. If None, no resize will be applied. 

29 """ 

30 

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

32 self.min_frame = kwargs.pop("min_frame", 0) 

33 self.colormap = kwargs.pop("colormap", "rgb") 

34 self.resize = kwargs.pop("resize", (DEFAULT_WIDTH, DEFAULT_HEIGHT)) 

35 self.interlaced = kwargs.pop("interlaced", False) 

36 

37 if self.min_frame < 0: 

38 self.min_frame = 0 

39 

40 if self.colormap not in ["rgb", "bgr"]: 

41 raise NotImplementedError 

42 

43 self.frame_reader = ControlledFPSVideoCapture(*args, **kwargs) 

44 

45 resize_log = f"Resize to {self.resize[0]}x{self.resize[1]}" if self.resize is not None else "No resize" 

46 colormap_log = f"Colormap: {self.colormap}" 

47 print(f"{resize_log}, {colormap_log}") 

48 

49 if self.min_frame > 0: 

50 self.frame_reader.set(cv2.CAP_PROP_POS_FRAMES, self.min_frame) 

51 

52 # According to our use-case we prefer to have the first frame already read. 

53 self.read() 

54 

55 def read(self): 

56 ret, self.image = self.frame_reader.read() 

57 

58 if self.interlaced: 

59 h, w = self.image.shape[:2] 

60 self.image = cv2.resize(self.image[::2], (w, h)) 

61 

62 self.image_before_resize = self.image 

63 if self.resize is not None: 

64 # ffmpeg prefers to save video with 1088 pixels height, thus we 

65 # add 8 pixels to our 1080p videos. The following removes those 

66 # pixels before resizing them to 720p images. 

67 if self.image.shape[0] == 1088: 

68 self.image = self.image[8:] 

69 

70 in_height, in_width = self.image.shape[:2] 

71 out_width, out_height = self.resize 

72 if in_height != out_height or in_width != out_width: 

73 in_aspect_ratio = float(in_height) / float(in_width) 

74 out_aspect_ratio = float(out_height) / float(out_width) 

75 if in_aspect_ratio != out_aspect_ratio: 

76 size_in = f"{in_width}x{in_height}" 

77 size_out = f"{out_width}x{out_height}" 

78 print("Warning: resizing image with a different aspect ratio") 

79 print(f" {size_in} -> {size_out}") 

80 self.image = cv2.resize(self.image, self.resize) 

81 if self.colormap == "rgb": 

82 self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB) 

83 

84 def set(self, prop_id, value): 

85 return self.frame_reader.set(prop_id, value) 

86 

87 def get(self, prop_id): 

88 return self.frame_reader.get(prop_id) 

89 

90 def seek(self, frame): 

91 """ 

92 This method aims to manage the tradoff between directly setting the 

93 VideoCapture to the desired frame and reading the successive frames 

94 up to the desired in order to get the frame in the fastest manner. 

95 

96 /!\\ Default value was chosen empirically and may not be optimal 

97 depending on your architecture. 

98 """ 

99 if not (self.frame_reader.output_frame <= frame < self.frame_reader.output_frame + SET_READ_TRADEOFF): 

100 self.set(cv2.CAP_PROP_POS_FRAMES, frame) 

101 self.read() 

102 

103 while self.frame_reader.output_frame < frame: 

104 self.read() 

105 return self.image 

106 

107 

108class StoredImagesVideoCapture(TransformedImageVideoCapture): 

109 """ 

110 This VideoCapture overlay stores previously read frames for faster access. 

111 

112 When setting to a new instant in the video, it prepares the previous 

113 frames by actually setting several frames back in order to facilitate 

114 backward reading. 

115 

116 The accessible frames are bounded as it was designed for the displayer and 

117 we may want a close up on a sequence (in particular for a better precision 

118 in the timeline). 

119 

120 Args: 

121 - min_frame: the index of the first accessible frame 

122 - max_frame: the index of the last accessible frame 

123 """ 

124 

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

126 self.stored_images = {} 

127 super().__init__(*args, **kwargs) 

128 

129 def get_max_frame(self): 

130 return super().get(cv2.CAP_PROP_FRAME_COUNT) 

131 

132 def read(self): 

133 """ 

134 Overrides read() method to store the output 

135 """ 

136 super().read() 

137 if self.frame_reader.output_frame - MAX_STORED in self.stored_images: 

138 del self.stored_images[self.frame_reader.output_frame - MAX_STORED] 

139 self.stored_images[self.frame_reader.output_frame] = self.image 

140 

141 def set(self, prop_id, value): 

142 """ 

143 Overrides set() method to erase all the stored values 

144 """ 

145 if prop_id == cv2.CAP_PROP_POS_FRAMES: 

146 self.stored_images = {} 

147 super().set(prop_id, value - AMONG_FRAMES) 

148 else: 

149 super().set(prop_id, value) 

150 

151 def seek(self, frame): 

152 if frame not in self.stored_images: 

153 super().seek(frame) 

154 self.frame = frame 

155 return self.stored_images[frame]