You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

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