From 858463255588df82d885e7576e1e39c6999ee7de Mon Sep 17 00:00:00 2001
From: j <j@mailb.org>
Date: Fri, 10 Aug 2018 09:11:40 +0100
Subject: [PATCH] followup to ba5e4af355b - create keyframe index as part of
 creating stream

---
 pandora/archive/chop.py    | 76 ++++++++++++++++++--------------------
 pandora/archive/extract.py |  4 +-
 pandora/archive/models.py  |  1 +
 3 files changed, 40 insertions(+), 41 deletions(-)

diff --git a/pandora/archive/chop.py b/pandora/archive/chop.py
index 3720bb28..9738c9b5 100644
--- a/pandora/archive/chop.py
+++ b/pandora/archive/chop.py
@@ -5,10 +5,39 @@ from bisect import bisect_left
 
 import ox
 
+def make_keyframe_index(video):
+    cmd = [
+        'ffprobe',
+        '-v', 'error',
+        '-show_packets', '-select_streams', 'v',
+        '-show_entries', 'packet=pts_time,flags',
+        '-of', 'csv',
+        '-i', video
+    ]
+    p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
+    stdout, stderr = p.communicate()
+    result = stdout.decode().strip()
+    keyframe_times = []
+    timecode = 0
+    for line in result.split('\n'):
+        if line.split(',')[1] != 'N/A':
+            timecode = line.split(',')[1]
+        if ',K' in line:
+            keyframe_times.append(float(timecode))
+
+    last_keyframe = ox.avinfo(video)['duration']
+    if keyframe_times[-1] != last_keyframe:
+        keyframe_times.append(last_keyframe)
+
+    keyframes_cache = video + '.keyframes'
+    with open(keyframes_cache, 'w') as fd:
+        json.dump(keyframe_times, fd, indent=0)
+    return keyframe_times
+
+
 class Chop(object):
     keyframes = []
     subtitles = None
-    info = {}
     ffmpeg = [
         'ffmpeg',
         '-nostats', '-loglevel', 'error',
@@ -77,12 +106,6 @@ class Chop(object):
         for segment in segments:
             os.unlink(segment)
 
-    def get_info(self):
-        if self.info:
-            return self.info
-        self.info = ox.avinfo(self.video)
-        return self.info
-
     def get_keyframes(self):
         video = self.video
         if self.keyframes:
@@ -93,35 +116,8 @@ class Chop(object):
             with open(keyframes_cache, 'r') as fd:
                 self.keyframes = json.load(fd)
             return self.keyframes
-
-        cmd = [
-            'ffprobe',
-            '-v', 'error',
-            '-show_packets', '-select_streams', 'v',
-            '-show_entries', 'packet=pts_time,flags',
-            '-of', 'csv',
-            '-i', video
-        ]
-        p = subprocess.Popen(cmd, stdout=subprocess.PIPE)
-        stdout, stderr = p.communicate()
-        result = stdout.decode().strip()
-        keyframe_times = []
-        timecode = 0
-        for line in result.split('\n'):
-            if line.split(',')[1] != 'N/A':
-                timecode = line.split(',')[1]
-            if ',K' in line:
-                keyframe_times.append(float(timecode))
-
-        last_keyframe = self.get_info()['duration']
-        if keyframe_times[-1] != last_keyframe:
-            keyframe_times.append(last_keyframe)
-
-        self.keyframes = keyframe_times
-        if not os.path.exists(keyframes_cache):
-            with open(keyframes_cache, 'w') as fd:
-                json.dump(keyframe_times, fd)
-        return keyframe_times
+        self.keyframes = make_keyframe_index(video)
+        return self.keyframes
 
     def get_gop_sections(self, start: float, end: float) -> dict:
         keyframes = self.get_keyframes()
@@ -141,17 +137,17 @@ class Chop(object):
         }
 
     def encode(self, source, target, start, duration):
-        info = self.get_info()
-        if self.info['audio']:
+        info = ox.avinfo(self.video)
+        if info['audio']:
             acodec = [
                 '-c:a',
-                self.info['audio'][0]['codec']
+                info['audio'][0]['codec']
             ]
         else:
             acodec = []
         vcodec = [
             '-c:v',
-            self.info['video'][0]['codec']
+            info['video'][0]['codec']
         ]
 
         cmd = self.ffmpeg + [
diff --git a/pandora/archive/extract.py b/pandora/archive/extract.py
index 41311812..f48ee9db 100644
--- a/pandora/archive/extract.py
+++ b/pandora/archive/extract.py
@@ -21,6 +21,8 @@ from ox.utils import json
 from django.conf import settings
 from PIL import Image
 
+from .chop import Chop, make_keyframe_index
+
 img_extension = 'jpg'
 
 MAX_DISTANCE = math.sqrt(3 * pow(255, 2))
@@ -347,6 +349,7 @@ def stream(video, target, profile, info, audio_track=0, flags={}):
         shutil.move(enc_target, target)
     for f in glob('%s.log*' % target):
         os.unlink(f)
+    make_keyframe_index(target)
     return True, None
 
 
@@ -614,7 +617,6 @@ def chop(video, start, end, subtitles=None):
     else:
         subtitles_f = None
     if ext == '.mp4':
-        from .chop import Chop
         Chop(video, choped_video, start, end, subtitles_f)
         if subtitles_f:
             os.unlink(subtitles_f)
diff --git a/pandora/archive/models.py b/pandora/archive/models.py
index ec0986e9..68ef30ab 100644
--- a/pandora/archive/models.py
+++ b/pandora/archive/models.py
@@ -370,6 +370,7 @@ class File(models.Model):
                         self.info.update(stream.info)
                         self.parse_info()
                         self.save()
+                    extract.make_keyframe_index(stream.media.path)
                 return True, stream.media.size
             return save_chunk(stream, stream.media, chunk, offset, name, done_cb)
         return False, 0