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

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 ctx1njm

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

31 self.url = url_vid 3 ctx1njm

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

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

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

35 

36 @property (empty)

37 def _url(self): (empty)

38 username = self._par.HTTP_DATA_SERVER_LOGIN 3 ctx1njm

39 password = self._par.HTTP_DATA_SERVER_PWD 3 ctx1njm

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

41 

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

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

44 return self.cap.read() 1 ctx1m

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 ctx1p

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

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

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

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

59 cap.release() 1 ctx1p

60 return frame_image 1 ctx1p

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)) 9 ctx1adebfchgi

83 if is_cap: 9 ctx1adebfchgi

84 cap = vid 

85 else: 

86 is_url = isinstance(vid, str) and vid.startswith('http') 9 ctx1adebfchgi

87 cap = VideoStreamer(vid).cap if is_url else cv2.VideoCapture(str(vid)) 9 ctx1adebfchgi

88 assert cap.isOpened(), 'Failed to open video' 9 ctx1adebfchgi

89 

90 frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) 9 ctx1adebfchgi

91 frame_numbers = frame_numbers if frame_numbers is not None else range(frame_count) 9 ctx1adebfchgi

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 9 ctx1adebfchgi

96 if as_list: 9 ctx1adebfchgi

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

98 else: 

99 ret, frame = cap.read() 9 ctx1adebfchgi

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

101 for ii, i in enumerate(frame_numbers): 9 ctx1adebfchgi

102 if not quiet: 9 ctx1adebfchgi

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

104 sys.stdout.flush() 9 ctx1adebfchgi

105 if to_set[ii]: 9 ctx1adebfchgi

106 cap.set(cv2.CAP_PROP_POS_FRAMES, i) 9 ctx1adebfchgi

107 ret, frame = cap.read() 9 ctx1adebfchgi

108 if ret: 9 ctx1adebfchgi

109 frame_images[ii] = func(frame[mask or ...]) 9 ctx1adebfchgi

110 else: 

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

112 if not is_cap: 9 ctx1adebfchgi

113 cap.release() 9 ctx1adebfchgi

114 if not quiet: 9 ctx1adebfchgi

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

116 return frame_images 9 ctx1adebfchgi

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') 4 ctx1abcj

128 cap = VideoStreamer(video_path).cap if is_url else cv2.VideoCapture(str(video_path)) 4 ctx1abcj

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

130 

131 # Get basic properties of video 

132 meta = Bunch() 4 ctx1abcj

133 meta.length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) 4 ctx1abcj

134 meta.fps = int(cap.get(cv2.CAP_PROP_FPS)) 4 ctx1abcj

135 meta.width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) 4 ctx1abcj

136 meta.height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) 4 ctx1abcj

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

138 if is_url and one: 4 ctx1abcj

139 eid = one.path2eid(video_path) 

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

141 label = label_from_path(video_path) 

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

143 assert len(record) == 1 

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

145 elif is_url and not one: 4 ctx1abcj

146 meta.size = None 

147 else: 

148 meta.size = Path(video_path).stat().st_size 4 ctx1abcj

149 cap.release() 4 ctx1abcj

150 return meta 4 ctx1abcj

151 

152 

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

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

155 

156 Parameters 

157 ---------- 

158 eid : UUID, str 

159 The session ID. 

160 label : str, tuple of str 

161 The video label (e.g. 'body') or a tuple thereof. 

162 one : one.api.One 

163 An instance of ONE. 

164 

165 Returns 

166 ------- 

167 str, dict of str 

168 The URL string if the label is a string, otherwise a dict of urls with labels as keys. 

169 

170 Raises 

171 ------ 

172 ValueError 

173 Video label is unreckognized. See `VIDEO_LABELS` for valid labels. 

174 """ 

175 valid_labels = VIDEO_LABELS 2 ctx1kl

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

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

178 one = one or ONE() 2 ctx1kl

179 session_path = one.eid2path(one.to_eid(eid)) 2 ctx1kl

180 

181 # Filter the video files 

182 def match(dataset): 2 ctx1kl

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

184 return matched and matched.group(1) in (label or valid_labels) 1 ctx1k

185 

186 if one.mode != 'remote': 2 ctx1kl

187 datasets = one.list_datasets(eid, details=False) 1 ctx1k

188 datasets = [ds for ds in datasets if match(ds)] 1 ctx1k

189 urls = [one.path2url(session_path / ds) for ds in datasets] 1 ctx1k

190 else: 

191 datasets = one.get_details(eid, full=True)['data_dataset_session_related'] 2 ctx1kl

192 urls = [ds['data_url'] for ds in datasets 2 ctx1kl

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

194 

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

196 if isinstance(label, str): 2 ctx1kl

197 return urls[0] 1 ctx1k

198 urls_dict = {label_from_path(url): url for url in urls} 2 ctx1kl

199 return {**dict.fromkeys(label), **urls_dict} if label else urls_dict 2 ctx1kl

200 

201 

202def label_from_path(video_name): (empty)

203 """ 

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

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

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

207 """ 

208 result = re.search(r'(?<=_)([a-z]+)(?=Camera)', str(video_name)) 5 ctx1aFkdG

209 return result.group() if result else None 5 ctx1aFkdG

210 

211 

212def assert_valid_label(label): (empty)

213 """ 

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

215 :param label: A video label to verify 

216 :return: the label in lowercase 

217 """ 

218 if not isinstance(label, str): 21 ctx1arostuqvwxdebyfzABCDE

219 try: 1 ctx1o

220 return tuple(map(assert_valid_label, label)) 1 ctx1o

221 except AttributeError: 1 ctx1o

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

223 if label.lower() not in VIDEO_LABELS: 21 ctx1arostuqvwxdebyfzABCDE

224 raise ValueError(f"camera must be one of ({', '.join(VIDEO_LABELS)})") 2 ctx1oq

225 return label.lower() 21 ctx1arostuqvwxdebyfzABCDE