Coverage for src/susi/analyse/movie_jitter.py: 27%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3 

4""" 

5Create movie incl. jitter and contrast plots 

6**Note:** requires ffmpeg executable! 

7 

8Author: feller 

9""" 

10 

11import numpy as np 

12import matplotlib.pyplot as plt 

13import os 

14import subprocess 

15from ..base import Logging 

16from ..io import FitsBatch 

17 

18logger = Logging.get_logger() 

19 

20DPI = 200 # frame resolution in dots per inch 

21FIGSIZE = (1920, 1080) # standard full-HD format 

22YLIM_SHIFTS = [-5, 5] # default plot range for pixel shifts 

23YLIM_CONTRAST = [0.3, 0.6] # default plot range for image contrast 

24COLORS = plt.rcParams['axes.prop_cycle'].by_key()['color'] # default matplotlib plotting color palette 

25 

26 

27def movie(config, batch_hk, sa, frame_range, roi): 

28 """ 

29 Create movie incl. jitter and contrast plots 

30 

31 ### Params 

32 - config: Config object 

33 - batch_hk: FitsBatch object with frame metadata 

34 - sa: result of ShiftAnalysis 

35 - frame_range: range of frames to include in the movie 

36 - roi: region of interest of the frames (YCEN, XCEN, YSIZE, XSIZE) 

37 """ 

38 video = __metadata(roi, frame_range) 

39 batch = __load_files(batch_hk, video) 

40 out_path = __create_out_path(config, video) 

41 

42 __render_frames(batch.data_array(), sa, video, out_path) 

43 __render_movie(config, video, out_path) 

44 

45 

46def __metadata(roi, frame_range): 

47 dy = round(roi[2] / 2) 

48 dx = round(roi[3] / 2) 

49 return { 

50 'name': 'movie', 

51 'range': range(frame_range[0], frame_range[1]), 

52 'roi': [slice(roi[0] - dy, roi[0] + dy), slice(roi[1] - dx, roi[1] + dx)] 

53 } 

54 

55 

56def __load_files(batch_hk, video): 

57 logger.info('Create list of filenames, timestamps') 

58 files = [batch_hk.file_by(i, full_path=True) for i in video['range']] 

59 

60 logger.info('Load image data from ROI %s', video['roi']) 

61 batch = FitsBatch(slices=video['roi']) 

62 batch.load(files) 

63 return batch 

64 

65 

66def __create_out_path(config, video): 

67 out_file = f"{config.data.dataset}_{config.cam.name}_{video['name']}_{video['range'].start}_{video['range'].stop}" 

68 out_path = os.path.join(config.data.out_path, out_file) 

69 logger.info('Output path: %s', out_path) 

70 if not os.path.isdir(out_path): 

71 logger.warning('Path does not exist, create path now') 

72 os.makedirs(out_path) 

73 return out_file 

74 

75 

76def __render_frames(data, sa, video, out_path): 

77 logger.info('Render movie frames') 

78 total_frames, contrast = __globals(data) 

79 fig, axs = plt.subplots(ncols=3, figsize=np.array(FIGSIZE) / DPI, dpi=DPI) 

80 for i in range(total_frames): 

81 __render_frame(axs, i, total_frames, sa, data, contrast, video) 

82 fig.tight_layout() 

83 fig.savefig(os.path.join(out_path, f"fig{i:05d}.png")) 

84 

85 

86def __globals(data): 

87 """compute number of frames and global data contrast""" 

88 return len(data), [np.std(im) / np.mean(im) for im in data] 

89 

90 

91def __render_frame(axs, fno, total_frames, sa, data, contrast, config): 

92 if (fno + 1) % 50 == 0: 

93 logger.debug('Frame no. %d of %d', fno + 1, total_frames) 

94 

95 im = data[fno, :] 

96 fnos = np.arange(config['range'].start, config['range'].start + fno + 1) 

97 

98 ax = axs[0] 

99 ax.imshow(im, cmap='gray') 

100 ax.set_title(f"Frame no. {fnos[fno]:05d}") 

101 ax.set_xlabel('Pixel') 

102 

103 ax = axs[1] 

104 ax.plot(fnos, sa['x'][fnos], color='blue') 

105 ax.plot(fnos, sa['y'][fnos], color='red') 

106 ax.set_xlim([config['range'].start, config['range'].stop]) 

107 ax.set_ylim([-5, 5]) 

108 ax.set_title('Shifts') 

109 ax.set_xlabel('Frame no.') 

110 ax.set_ylabel('Pixels') 

111 

112 ax = axs[2] 

113 ax.plot(fnos, contrast[0:fno + 1], color='blue') 

114 ax.set_xlim([config['range'].start, config['range'].stop]) 

115 ax.set_ylim([0.3, 0.6]) 

116 ax.set_title('Contrast') 

117 ax.set_xlabel('Frame no.') 

118 

119 

120def __render_movie(config, video, out_path): 

121 logger.info('Render movie') 

122 movie_name = f"{config.data.dataset}_{config.cam.name}_{video['name']}" 

123 movie_name += f"_{video['range'].start}_{video['range'].stop}.mp4" 

124 args = ['ffmpeg', '-r', '10', '-f', 'image2', '-i', out_path + '/fig%05d.png', '-vcodec', 'libx264', 

125 '-crf', '25', '-pix_fmt', 'yuv420p', movie_name] 

126 subprocess.run(args)