Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.

166 řádky
6.3 KiB

  1. """Classes for producing videos"""
  2. from abc import ABC
  3. import json
  4. import logging
  5. import os
  6. import subprocess
  7. import tempfile
  8. # for visualisations:
  9. import matplotlib.pyplot as plt
  10. # for encoding as JSON
  11. from pipeline.utils import Feature
  12. class Producer(ABC):
  13. """Generic producer interface."""
  14. def __init__(self, features):
  15. """All producers should take a list of features as input"""
  16. def produce(self):
  17. """All Producers should produce something!"""
  18. class VideoProducer(Producer):
  19. """Video producer interface."""
  20. class FfmpegVideoProducer(VideoProducer):
  21. """Produce videos using ffmpeg"""
  22. # TODO: consider output filename options
  23. def __init__(self, features):
  24. if not features:
  25. raise ValueError("No features provided")
  26. # TODO: consider if we want to permit empty features (producing no video)
  27. self.features = features
  28. def _ffmpeg_feature_to_clip(self, feature=None, output_filepath=None):
  29. """use ffmpeg to produve a video clip from a feature"""
  30. OVERWRITE = True # TODO: consider making this a config option
  31. if not feature or not feature.interval:
  32. raise ValueError("No feature provided")
  33. if not output_filepath:
  34. raise ValueError("No output filepath provided")
  35. ffmpeg_prefix = ["ffmpeg", "-y"] if OVERWRITE else ["ffmpeg"]
  36. ffmpeg_suffix = ["-r", "60", "-c:v", "libx264", "-crf", "26", "-c:a", "aac", "-preset", "ultrafast"]
  37. # TODO: match framerate of input video
  38. # TODO: adjustable encoding options
  39. seek = ["-ss", str(feature.interval.start)]
  40. duration = ["-t", str(feature.interval.duration)]
  41. ffmpeg_args = ffmpeg_prefix + seek + ["-i"] + [feature.path] + duration + ffmpeg_suffix + [output_filepath]
  42. logging.info(f"ffmpeg_args: {ffmpeg_args}")
  43. subprocess.run(ffmpeg_args, stdout=None, stderr=None)
  44. def _ffmpeg_concat_clips(self, clips=None, output_filepath=None):
  45. """use ffmpeg to concatenate clips into a single video"""
  46. OVERWRITE = True
  47. ffmpeg_prefix = ["ffmpeg"]
  48. ffmpeg_prefix += ["-y"] if OVERWRITE else []
  49. ffmpeg_prefix += ["-f", "concat", "-safe", "0", "-i"]
  50. # there is a method to do this via process substitution, but it's not portable
  51. # so we'll use the input file list method
  52. if not clips:
  53. raise ValueError("No clips provided")
  54. if not output_filepath:
  55. raise ValueError("No output filepath provided")
  56. # generate a temporary file with the list of clips
  57. join_file = tempfile.NamedTemporaryFile(mode="w")
  58. for clip in clips:
  59. join_file.write(f"file '{clip}'\n")
  60. join_file.flush()
  61. ffmpeg_args = ffmpeg_prefix + [join_file.name] + ["-c", "copy", output_filepath]
  62. logging.info(f"ffmpeg_args: {ffmpeg_args}")
  63. subprocess.run(ffmpeg_args, stdout=None, stderr=None)
  64. join_file.close()
  65. def produce(self):
  66. OUTPUT_DIR = "/tmp/" # TODO: make this a config option
  67. clips = []
  68. for num, feature in enumerate(self.features):
  69. output_filepath = f"{OUTPUT_DIR}/highlight_{num}.mp4"
  70. self._ffmpeg_feature_to_clip(feature, output_filepath)
  71. clips.append(output_filepath)
  72. # concatenate the clips
  73. output_filepath = f"{OUTPUT_DIR}/highlights.mp4"
  74. self._ffmpeg_concat_clips(clips, output_filepath)
  75. logging.info(f"Produced video: {output_filepath}")
  76. class VisualisationProducer(Producer):
  77. """Visualisation producer -- illustrate the features we have extracted"""
  78. def __init__(self, features):
  79. if not features:
  80. raise ValueError("No features provided")
  81. self.features = features
  82. def produce(self):
  83. """Produce visualisation"""
  84. # basic idea: use matplotlib to plot:
  85. # - a wide line segment representing the source video[s]
  86. # - shorter line segments representing the features extracted where:
  87. # + width represents duration
  88. # + colour represents feature type
  89. # + position represents time
  90. # - save as image
  91. plotted_source_videos = []
  92. bar_labels = []
  93. fig, ax = plt.subplots()
  94. for feature in self.features:
  95. # plot source video line if not done already
  96. if feature.source not in plotted_source_videos:
  97. # use video duration as width
  98. # ax.plot([0, feature.source.duration()], [0, 0], color='black', linewidth=10)
  99. ax.broken_barh([(0, feature.source.duration())], (0, 5), facecolors='grey')
  100. plotted_source_videos.append(feature.source)
  101. bar_labels.append(os.path.basename(feature.source.path))
  102. # annotate the source video
  103. ax.text(0.25, 0.25, os.path.basename(feature.source.path), ha='left', va='bottom',
  104. fontsize=16)
  105. # plot feature line
  106. # ax.plot([feature.interval.start, feature.interval.end], [1, 1], color='red', linewidth=5)
  107. ax.broken_barh([(feature.interval.start, feature.interval.duration)], (10, 5), facecolors='red')
  108. if feature.feature_extractor not in bar_labels:
  109. bar_labels.append(feature.feature_extractor)
  110. # label bar with feature extractor
  111. ax.text(0, 8, feature.feature_extractor, ha='left', va='bottom',
  112. fontsize=16)
  113. # label the plot's axes
  114. ax.set_xlabel('Time')
  115. # ax.set_yticks([], labels=bar_labels)
  116. ax.set_yticks([])
  117. # ax.tick_params(axis='y', labelrotation=90, ha='right')
  118. # save the plot
  119. plt.savefig("/tmp/visualisation.png")
  120. plt.close()
  121. class PipelineJSONEncoder(json.JSONEncoder):
  122. def default(self, obj):
  123. if hasattr(obj, 'to_json'):
  124. return obj.to_json()
  125. else:
  126. return json.JSONEncoder.default(self, obj)
  127. class JSONProducer(Producer):
  128. """Produce JSON output"""
  129. def __init__(self, features):
  130. if not features:
  131. raise ValueError("No features provided")
  132. self.features = features
  133. def produce(self):
  134. # FIXME: config option for output path
  135. with open("/tmp/features.json", "w") as jsonfile:
  136. jsonfile.write(json.dumps(self.features, cls=PipelineJSONEncoder, indent=4))