Non puoi selezionare più di 25 argomenti Gli argomenti devono iniziare con una lettera o un numero, possono includere trattini ('-') e possono essere lunghi fino a 35 caratteri.

producers.py 6.5 KiB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173
  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))