improved multithreading

Significantly reduced the memory usage

Author: private-yusuke

t2v.py

@@ -34,9 +34,8 @@ def getDisplayMode(mode):
 
 
 def textToFrame(data):
-    ind = data[0]
-    text = data[1]
-    opts = data[2]
+    text = data[0]
+    opts = data[1]
     # print(os.getpid())
 
     monospace = cv2.freetype.createFreeType2()
@@ -59,7 +58,7 @@ def textToFrame(data):
     frame = cv2.resize(frame, dsize=(
         opts['output_width'], opts['output_height']))
     # cv2.imwrite('test/{:04}.png'.format(ind), frame)
-    return (ind, frame)
+    return frame
 
 
 def main():
@@ -105,24 +104,15 @@ def main():
             output = args.output + '_tmp.mp4'
         print(output)
 
-        frames_unord = []
-        with mp.Pool(mp.cpu_count()) as pool:
-            with tqdm(total=len(texts)) as t:
-                for res in pool.imap_unordered(
-                        textToFrame, zip(range(len(texts)), texts, repeat(opts))):
-                    t.update()
-                    frames_unord.append(res)
-
-        frames = sorted(frames_unord)
-        frames = list(map(lambda p: p[1], frames))
-
         out = cv2.VideoWriter(output, cv2.VideoWriter_fourcc(
             *'mp4v'), FPS, (WIDTH, HEIGHT))
 
-        with tqdm(total=len(frames)) as t:
-            for frame in frames:
-                out.write(frame)
-                t.update()
+        with mp.Pool(mp.cpu_count()) as pool:
+            with tqdm(total=len(texts)) as t:
+                for res in pool.imap(
+                        textToFrame, zip(texts, repeat(opts))):
+                    t.update()
+                    out.write(res)
 
         out.release()
 

v2t.py

@@ -54,10 +54,7 @@ def slicer(a, chunk_i, chunk_j, two_d=True):
     return c
 
 
-def frameToText(data):
-    k = data[0]
-    frame = data[1]
-    opts = data[2]
+def frameToText(frame, opts):
     text = ''
 
     frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
@@ -87,13 +84,17 @@ def frameToText(data):
     tmp3 = [tmp2[i:i+chunk_size] for i in range(0, len(tmp2), chunk_size)]
 
     text += '\n'.join(tmp3)
-    return (k, text)
+    return text
 
 
-def loadFrame(data):
-    ind = data[0]
-    path = data[1]
-    return (ind, cv2.imread(path))
+def loadFrame(path):
+    return cv2.imread(path)
+
+
+def loadFrameAndConvertToText(data):
+    path = data[0]
+    opts = data[1]
+    return frameToText(loadFrame(path), opts)
 
 
 def main(args):
@@ -120,38 +121,24 @@ def main(args):
     opts['PALETTE'] = PALETTE
     opts['mode'] = args.mode
 
-    frames_unord = []
-
     with tempfile.TemporaryDirectory() as tmpdir:
         print(tmpdir)
         ffmpeg.input(filepath).output(os.path.join(tmpdir, '%d.png')).run()
         files = os.listdir(tmpdir)
         files.sort(key=lambda f: int(re.sub('\D', '', f)))
         with mp.Pool(mp.cpu_count()) as pool:
             with tqdm(total=len(files)) as t:
-                for res in pool.imap_unordered(loadFrame, zip(range(len(files)), list(map(lambda f: os.path.join(tmpdir, f), files)))):
-                    frames_unord.append(res)
+                for res in pool.imap(loadFrameAndConvertToText, zip(list(map(lambda f: os.path.join(tmpdir, f), files)), repeat(opts))):
+                    texts.append(res)
                     t.update()
 
-    frames = sorted(frames_unord)
-    frames = list(map(lambda p: p[1], frames))
-
-    texts_unord = []
-
-    with mp.Pool(mp.cpu_count()) as pool:
-        with tqdm(total=len(frames)) as t:
-            for res in pool.imap_unordered(frameToText, zip(range(len(frames)), frames, repeat(opts))):
-                t.update()
-                texts_unord.append(res)
-
-    texts = sorted(texts_unord)
-    texts = list(map(lambda p: p[1], texts))
-
     print('frames loaded')
 
-    fps = MediaInfo.parse(filepath).tracks[0].frame_rate
-    origwidth = frames[0].shape[1]
-    origheight = frames[0].shape[0]
+    media_info = [info for info in MediaInfo.parse(
+        filepath).tracks if info.track_type == 'Video'][0]
+    fps = media_info.frame_rate
+    origwidth = media_info.width
+    origheight = media_info.height
     width = ceil(origwidth / N_WIDTH)
     height = ceil(origheight / N_HEIGHT)
 
@@ -160,9 +147,7 @@ def main(args):
         f.write("{},{},{},{},{},{}\n".format(
             filepath, fps, int(width), int(height), origwidth, origheight))
         f.write('=====\n')
-        for text in texts:
-            f.write(text)
-            f.write('\n')
+        f.write('\n'.join(texts))
 
 
 if __name__ == "__main__":