Advertisement
Guest User

Untitled

a guest
Feb 10th, 2021
119
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 7.63 KB | None | 0 0
  1. from PIL import Image
  2. from PIL import ImageStat
  3. import numpy as np
  4. import matplotlib.pyplot as plt
  5. from scipy.interpolate import PchipInterpolator
  6. from cv2 import cv2
  7. import subprocess
  8.  
  9. ffmpeg_executable = 'ffmpeg'
  10.  
  11. def fps_to_ffmpeg_fps(fps: float):
  12. if round(fps * 100) == 5994:
  13. return '60000/1001'
  14. elif round(fps * 100) % 100 == 0:
  15. return str(int(fps))
  16. return str(fps)
  17.  
  18. def dnxhd_bitrate(w: int, h: int, fps: float):
  19. if w == 1920 and h == 1080:
  20. if int(round(fps)) == 60:
  21. return 'yuv422p', '440M', None
  22. elif int(round(fps)) == 50:
  23. return 'yuv422p', '365M', None
  24. elif int(round(fps)) == 30:
  25. return 'yuv422p', '220M', None
  26. elif w == 3840 and h == 2160:
  27. if int(round(fps)) == 60:
  28. return 'yuv422p10le', '1760M', 'dnxhr_hqx'
  29. elif int(round(fps)) == 50:
  30. return 'yuv422p10le', '1460M', 'dnxhr_hqx'
  31. elif int(round(fps)) == 30:
  32. return 'yuv422p10le', '880M', 'dnxhr_hqx'
  33. raise Exception(f'fps {fps} incompatible with dnxhd')
  34.  
  35. class FfmpegVideoWriter:
  36. def __init__(self, path: str, quality: int, w: int, h: int, duration_s: float, fps: float, filters: str = None, preset: str = 'slow', audio_file: str = None):
  37. assert quality >= 0 and quality < 51
  38. fps_str = format(fps, '.2f')
  39. args = [ffmpeg_executable, '-y', '-f', 'image2pipe', '-vcodec', 'bmp', '-t', '00:{0}:{1}'.format(int(duration_s // 60), duration_s % 60), '-video_size', '{0}x{1}'.format(int(w), int(h)) , '-r', fps_str, '-i', '-']
  40. if audio_file:
  41. args += ['-i', audio_file, '-map', '0:v:0', '-map', '1:a:0']
  42. if filters:
  43. args += ['-vf', filters]
  44.  
  45. res_str = f'{w}x{h}'
  46. pix_fmt, bitrate, profile = dnxhd_bitrate(w, h, fps)
  47. args += ['-vcodec', 'dnxhd', '-acodec', 'pcm_s16le', '-s', res_str, '-b:v', bitrate, '-s', res_str, '-r', fps_to_ffmpeg_fps(fps), '-pix_fmt', pix_fmt]
  48. if profile:
  49. args += ['-profile:v', profile]
  50. args += ['-f', 'mov', path]
  51. # args += ['-vcodec', 'dnxhd', '-acodec', 'pcm_s16le', '-s', '1920x1080', '-b:v', '365M', '-s', '1920x1080', '-r', '50', '-pix_fmt', 'yuv422p', '-f', 'mov', path]
  52. # args += ['-vcodec', 'libx264', '-crf', str(quality), '-preset', preset, '-r', fps_str, path]
  53. print(args)
  54. self.p = subprocess.Popen(args, stdin=subprocess.PIPE)
  55.  
  56. def write_frame(self, cvframe):
  57. rgbimg = cv2.cvtColor(np.array(cvframe, dtype=np.uint8), cv2.COLOR_BGR2RGB)
  58. img = Image.fromarray(rgbimg)
  59. img.save(self.p.stdin, 'BMP')
  60.  
  61. def release(self):
  62. self.p.stdin.close()
  63. self.p.wait()
  64.  
  65. def create_curve(hist, cutpoint: float, cutpoint2: float):
  66. pixel_count = np.sum(hist)
  67. cumhist = np.cumsum(hist)
  68. thresh = pixel_count * cutpoint
  69. spointx = np.sum(np.less(cumhist, thresh))
  70.  
  71. if spointx / 255 >= cutpoint:
  72. return np.arange(len(hist))
  73.  
  74. # print(cutpoint, spointx, max(hist) / pixel_count / (spointx / 255))
  75. x = np.arange(len(hist))
  76. # cs = PchipInterpolator([0, spointx * 0.8, spointx, 255], [0, cutpoint2 * 0.9, cutpoint2, 1])
  77. # cs = PchipInterpolator([0, spointx, 255], [0, cutpoint2, 1])
  78.  
  79.  
  80. # Если много льда, то мы хотим его сжать, чем менее равномерно, тем более
  81. peak = max(hist) / pixel_count / (spointx / 255)
  82. ice_compression = (max(min(peak, 0.1), 0.01) - 0.01)
  83. # print(ice_compression)
  84.  
  85. xs = [0, spointx * (0.9 - ice_compression), spointx, 255]
  86. ys = [0, cutpoint2 * 0.9, cutpoint2, 1]
  87.  
  88. cs = PchipInterpolator(xs, ys)
  89.  
  90. # Рисуем
  91. # fig, ax = plt.subplots(figsize=(6.5, 4))
  92. # ax.plot(x, hist / pixel_count * 50)
  93. # ax.plot(x, cs(x))
  94. # ax.plot(xs, ys, 'o')
  95. # plt.show()
  96.  
  97. # print(cumhist, pixel_count * cutpoint)
  98. return np.round(cs(x) * 255)
  99.  
  100. def make_monotonic(a):
  101. amon = [a[0]]
  102. last = a[0]
  103.  
  104. for x in a[1:]:
  105. if x >= last:
  106. amon.append(x)
  107. last = x
  108. else:
  109. amon.append(last)
  110.  
  111. return np.array(amon, dtype='uint8')
  112.  
  113. def pic1(im):
  114. # im = Image.open("/media/vlad/new2021/editins/sasha/jumpfest/day3/mpv-shot0002.jpg")
  115.  
  116. rgb_hist = np.array(im.histogram()[:256]) + np.array(im.histogram()[256:512]) + np.array(im.histogram()[512:768])
  117. crv = create_curve(np.array(rgb_hist), 0.99, 0.96)
  118. crvmon = np.array(make_monotonic(list(crv)), dtype='uint8')
  119. print(len(crvmon))
  120.  
  121. im2 = im.point(lambda x: crvmon[x])
  122. return im2
  123.  
  124. class BufferedVideoInput:
  125. def __init__(self, vid, lookahead, fun):
  126. self.vid = vid
  127. self.cached_frames = []
  128. self.lookahead = lookahead
  129. self.fun = fun
  130. self.cached_fun_results = []
  131. self.pos = -1
  132. self.eof = False
  133.  
  134. def get_fun_results(self, offset):
  135. idx = self.pos + offset - 1
  136. return self.cached_fun_results[max(min(idx, len(self.cached_fun_results) - 1), 0)]
  137.  
  138. def read(self):
  139. while not self.eof and len(self.cached_frames) < self.lookahead:
  140. retval, frame = self.vid.read()
  141. if not retval:
  142. self.eof = True
  143. else:
  144. self.cached_fun_results.append(self.fun(frame))
  145. self.cached_frames.append(frame)
  146.  
  147. self.pos += 1
  148. if len(self.cached_frames) == 0:
  149. return False, None
  150. else:
  151. cached_frame = self.cached_frames[0]
  152. self.cached_frames.pop(0)
  153. return True, cached_frame
  154.  
  155. def get(self, cap):
  156. return self.vid.get(cap)
  157.  
  158. def release(self):
  159. self.vid.release()
  160.  
  161. def create_brightness_from_video_capture(vid_capture):
  162. def make_curve(frame):
  163. hist_r = cv2.calcHist([frame], [0], None, [256], [0,256])
  164. hist_g = cv2.calcHist([frame], [1], None, [256], [0,256])
  165. hist_b = cv2.calcHist([frame], [2], None, [256], [0,256])
  166. rgb_hist = np.array(hist_r + hist_g + hist_b)
  167. crv = create_curve(np.array(rgb_hist), 0.99, 0.96)
  168. crvmon = make_monotonic(list(crv))
  169. return crvmon
  170.  
  171. return BufferedVideoInput(vid_capture, 15, make_curve)
  172.  
  173. def vid_bright(input_path, output_path):
  174. vid = create_brightness_from_video_capture(cv2.VideoCapture(input_path))
  175.  
  176. input_width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
  177. input_height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
  178. input_fps = vid.get(cv2.CAP_PROP_FPS)
  179. framecount = vid.get(cv2.CAP_PROP_FRAME_COUNT)
  180. filters_str = ''
  181. quality = 12
  182.  
  183. print(input_fps)
  184. print(filters_str)
  185.  
  186. outputvid = FfmpegVideoWriter(output_path, fps=input_fps, quality=quality, filters=filters_str, w=input_width, h=input_height, duration_s=framecount / input_fps, audio_file=input_path)
  187.  
  188. i = 1
  189. while True:
  190. retval, frame = vid.read()
  191.  
  192. if not retval:
  193. break
  194.  
  195. # Скользящее окно с ядром [1, 1, .., 1]
  196. curve = vid.get_fun_results(1)
  197. total_curve = np.zeros(curve.shape)
  198. total_curve += vid.get_fun_results(1)
  199. for i in range(5):
  200. total_curve += vid.get_fun_results(1 - i)
  201. total_curve += vid.get_fun_results(1 + i)
  202. total_curve = np.array(total_curve / (5 * 2 + 1), dtype='uint8')
  203.  
  204. frame_cc = total_curve[frame]
  205.  
  206. outputvid.write_frame(frame_cc)
  207. i += 1
  208.  
  209.  
  210. outputvid.release()
  211. vid.release()
  212.  
  213.  
  214. path_list = []
  215.  
  216.  
  217.  
  218. for path in path_list:
  219. vid_bright(path, path + '.bright.mp4')
  220.  
  221.  
  222. # im.show()
  223. # im2.show()
  224.  
  225.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement