Coverage for ibllib/io/video.py: 97%

1"""Functions for fetching video frames, meta data and file locations""" 

2import sys (empty)

3import re (empty)

4from datetime import timedelta (empty)

5from pathlib import Path (empty)

6 

7import cv2 (empty)

8import numpy as np (empty)

9 

10from iblutil.util import Bunch (empty)

11from one.api import ONE (empty)

12from one import params (empty)

13 

14VIDEO_LABELS = ('left', 'right', 'body', 'belly') (empty)

15 

16 

17class VideoStreamer: (empty)

18 """ 

19 Provides a wrapper to stream a video from a password protected HTTP server using opencv 

20 """ 

21 

22 def __init__(self, url_vid): (empty)

23 """ 

24 TODO Allow auth as input 

25 :param url_vid: full url of the video or dataset dictionary as output by alyx rest datasets 

26 :returns cv2.VideoCapture object 

27 """ 

28 # pop the data url from the dataset record if the input is a dictionary 

29 if isinstance(url_vid, dict): 3 ctx1vgu

30 url_vid = next(fr['data_url'] for fr in url_vid['file_records'] if fr['data_url']) 1 ctx1u

31 self.url = url_vid 3 ctx1vgu

32 self._par = params.get(silent=True) 3 ctx1vgu

33 self.cap = cv2.VideoCapture(self._url) 3 ctx1vgu

34 self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT)) 3 ctx1vgu

35 

36 @property (empty)

37 def _url(self): (empty)

38 username = self._par.HTTP_DATA_SERVER_LOGIN 3 ctx1vgu

39 password = self._par.HTTP_DATA_SERVER_PWD 3 ctx1vgu

40 return re.sub(r'(^https?://)', r'\1' + f'{username}:{password}@', self.url) 3 ctx1vgu

41 

42 def get_frame(self, frame_index): (empty)

43 self.cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index) 1 ctx1u

44 return self.cap.read() 1 ctx1u

45 

46 

47def get_video_frame(video_path, frame_number): (empty)

48 """ 

49 Obtain numpy array corresponding to a particular video frame in video_path 

50 :param video_path: local path to mp4 file 

51 :param frame_number: video frame to be returned 

52 :return: numpy array corresponding to frame of interest. Dimensions are (w, h, 3) 

53 """ 

54 is_url = isinstance(video_path, str) and video_path.startswith('http') 1 ctx1A

55 cap = VideoStreamer(video_path).cap if is_url else cv2.VideoCapture(str(video_path)) 1 ctx1A

56 # 0-based index of the frame to be decoded/captured next. 

57 cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number) 1 ctx1A

58 ret, frame_image = cap.read() 1 ctx1A

59 cap.release() 1 ctx1A

60 return frame_image 1 ctx1A

61 

62 

63def get_video_frames_preload(vid, frame_numbers=None, mask=Ellipsis, as_list=False, (empty)

64 func=lambda x: x, quiet=False): 

65 """ 

66 Obtain numpy array corresponding to a particular video frame in video. 

67 Fetching and returning a list is about 33% faster but may be less memory controlled. NB: Any 

68 gain in speed will be lost if subsequently converted to array. 

69 :param vid: URL or local path to mp4 file or cv2.VideoCapture instance. 

70 :param frame_numbers: video frames to be returned. If None, return all frames. 

71 :param mask: a logical mask or slice to apply to frames 

72 :param as_list: if true the frames are returned as a list, this is faster but may be less 

73 memory efficient 

74 :param func: Function to be applied to each frame. Applied after masking if applicable. 

75 :param quiet: if true, suppress frame loading progress output. 

76 :return: numpy array corresponding to frame of interest, or list if as_list is True. 

77 Default dimensions are (n, w, h, 3) where n = len(frame_numbers) 

78 

79 Example - Load first 1000 frames, keeping only the first colour channel: 

80 frames = get_video_frames_preload(vid, range(1000), mask=np.s_[:, :, 0]) 

81 """ 

82 is_cap = not isinstance(vid, (str, Path)) 11 ctx1abfhdiekjlc

83 if is_cap: 11 ctx1abfhdiekjlc

84 cap = vid 

85 else: 

86 is_url = isinstance(vid, str) and vid.startswith('http') 11 ctx1abfhdiekjlc

87 cap = VideoStreamer(vid).cap if is_url else cv2.VideoCapture(str(vid)) 11 ctx1abfhdiekjlc

88 assert cap.isOpened(), 'Failed to open video' 11 ctx1abfhdiekjlc

89 

90 frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) 11 ctx1abfhdiekjlc

91 frame_numbers = frame_numbers if frame_numbers is not None else range(frame_count) 11 ctx1abfhdiekjlc

92 

93 # Setting the index is extremely slow; determine where frame index must be set 

94 # The first index is always explicitly set. 

95 to_set = np.insert(np.diff(frame_numbers), 0, 0) != 1 11 ctx1abfhdiekjlc

96 if as_list: 11 ctx1abfhdiekjlc

97 frame_images = [None] * len(frame_numbers) 1 ctx1j

98 else: 

99 ret, frame = cap.read() 11 ctx1abfhdiekjlc

100 frame_images = np.zeros((len(frame_numbers), *func(frame[mask or ...]).shape), np.uint8) 10 ctx1abhdiekjlc

101 for ii, i in enumerate(frame_numbers): 10 ctx1abhdiekjlc

102 if not quiet: 10 ctx1abhdiekjlc

103 sys.stdout.write(f'\rloading frame {ii}/{len(frame_numbers)}') 10 ctx1abhdiekjlc

104 sys.stdout.flush() 10 ctx1abhdiekjlc

105 if to_set[ii]: 10 ctx1abhdiekjlc

106 cap.set(cv2.CAP_PROP_POS_FRAMES, i) 10 ctx1abhdiekjlc

107 ret, frame = cap.read() 10 ctx1abhdiekjlc

108 if ret: 10 ctx1abhdiekjlc

109 frame_images[ii] = func(frame[mask or ...]) 10 ctx1abhdiekjlc

110 else: 

111 print(f'failed to read frame #{i}') 

112 if not is_cap: 10 ctx1abhdiekjlc

113 cap.release() 10 ctx1abhdiekjlc

114 if not quiet: 10 ctx1abhdiekjlc

115 sys.stdout.write('\x1b[2K\r') # Erase current line in stdout 10 ctx1abhdiekjlc

116 return frame_images 10 ctx1abhdiekjlc

117 

118 

119def get_video_meta(video_path, one=None): (empty)

120 """ 

121 Return a bunch of video information with the fields ('length', 'fps', 'width', 'height', 

122 'duration', 'size') 

123 :param video_path: A path to the video. May be a file path or URL. 

124 :param one: An instance of ONE 

125 :return: A Bunch of video mata data 

126 """ 

127 is_url = isinstance(video_path, str) and video_path.startswith('http') 13 ctx1abopqrfdensgc

128 cap = VideoStreamer(video_path).cap if is_url else cv2.VideoCapture(str(video_path)) 13 ctx1abopqrfdensgc

129 assert cap.isOpened(), f'Failed to open video file {video_path}' 13 ctx1abopqrfdensgc

130 

131 # Get basic properties of video 

132 meta = Bunch() 13 ctx1abopqrfdensgc

133 meta.length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) 13 ctx1abopqrfdensgc

134 meta.fps = int(cap.get(cv2.CAP_PROP_FPS)) 13 ctx1abopqrfdensgc

135 meta.width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) 13 ctx1abopqrfdensgc

136 meta.height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) 13 ctx1abopqrfdensgc

137 meta.duration = timedelta(seconds=meta.length / meta.fps) if meta.fps > 0 else 0 13 ctx1abopqrfdensgc

138 if is_url and one: 13 ctx1abopqrfdensgc

139 eid = one.path2eid(video_path) 1 ctx1g

140 datasets = one.list_datasets(eid, details=True) 1 ctx1g

141 label = label_from_path(video_path) 1 ctx1g

142 record = datasets[datasets['rel_path'].str.contains(f'_iblrig_{label}Camera.raw')] 1 ctx1g

143 assert len(record) == 1 1 ctx1g

144 meta.size = record['file_size'].iloc[0] 1 ctx1g

145 elif is_url and not one: 13 ctx1abopqrfdensgc

146 meta.size = None 

147 else: 

148 meta.size = Path(video_path).stat().st_size 13 ctx1abopqrfdensgc

149 cap.release() 13 ctx1abopqrfdensgc

150 return meta 13 ctx1abopqrfdensgc

151 

152 

153def url_from_eid(eid, label=None, one=None): (empty)

154 """Return the video URL(s) for a given eid 

155 

156 :param eid: The session id 

157 :param label: The video label (e.g. 'body') or a tuple thereof 

158 :param one: An instance of ONE 

159 :return: The URL string if the label is a string, otherwise a dict of urls with labels as keys 

160 """ 

161 valid_labels = VIDEO_LABELS 2 ctx1mt

162 if not (label is None or np.isin(label, valid_labels).all()): 2 ctx1mt

163 raise ValueError('labels must be one of ("%s")' % '", "'.join(valid_labels)) 1 ctx1m

164 one = one or ONE() 2 ctx1mt

165 session_path = one.eid2path(one.to_eid(eid)) 2 ctx1mt

166 

167 # Filter the video files 

168 def match(dataset): 2 ctx1mt

169 matched = re.match(r'(?:_iblrig_)([a-z]+)(?=Camera.raw.mp4$)', dataset.rsplit('/')[-1]) 2 ctx1mt

170 return matched and matched.group(1) in (label or valid_labels) 2 ctx1mt

171 

172 if one.mode != 'remote': 2 ctx1mt

173 datasets = one.list_datasets(eid, details=False) 2 ctx1mt

174 datasets = [ds for ds in datasets if match(ds)] 2 ctx1mt

175 urls = [one.path2url(session_path / ds) for ds in datasets] 2 ctx1mt

176 else: 

177 datasets = one.get_details(eid, full=True)['data_dataset_session_related'] 1 ctx1m

178 urls = [ds['data_url'] for ds in datasets 1 ctx1m

179 if ds['dataset_type'] == '_iblrig_Camera.raw' and match(ds['name'])] 

180 

181 # If one label specified, return the url, otherwise return a dict 

182 if isinstance(label, str): 2 ctx1mt

183 return urls[0] 1 ctx1m

184 urls_dict = {label_from_path(url): url for url in urls} 2 ctx1mt

185 return {**dict.fromkeys(label), **urls_dict} if label else urls_dict 2 ctx1mt

186 

187 

188def label_from_path(video_name): (empty)

189 """ 

190 Return the video label, e.g.. 'left', 'right' or 'body' 

191 :param video_name: A file path, URL or file name for the video 

192 :return: The string label or None if the video doesn't match 

193 """ 

194 result = re.search(r'(?<=_)([a-z]+)(?=Camera)', str(video_name)) 16 ctx1aRmbSBwxfhTUVngc

195 return result.group() if result else None 16 ctx1aRmbSBwxfhTUVngc

196 

197 

198def assert_valid_label(label): (empty)

199 """ 

200 Raises a value error is the provided label is not supported. 

201 :param label: A video label to verify 

202 :return: the label in lowercase 

203 """ 

204 if not isinstance(label, str): 27 ctx1aEyFGHDIbBwxfJhdKiLzCMNOPQc

205 try: 6 ctx1ywxfzC

206 return tuple(map(assert_valid_label, label)) 6 ctx1ywxfzC

207 except AttributeError: 2 ctx1yz

208 raise ValueError('label must be string or iterable of strings') 

209 if label.lower() not in VIDEO_LABELS: 27 ctx1aEyFGHDIbBwxfJhdKiLzCMNOPQc

210 raise ValueError(f"camera must be one of ({', '.join(VIDEO_LABELS)})") 3 ctx1yDz

211 return label.lower() 27 ctx1aEyFGHDIbBwxfJhdKiLzCMNOPQc