interactive converter now saves the session.

if input frames are changed (amount or filenames)
then interactive converter automatically starts a new session.
if model is more trained then all frames will be recomputed again with their saved configs.

interact/interact.py

@@ -137,9 +137,9 @@ class InteractBase(object):
             else: print("capture_keys: already set for window ", wnd_name)
         else: print("capture_keys: named_window ", wnd_name, " not found.")
 
-    def progress_bar(self, desc, total, leave=True):
+    def progress_bar(self, desc, total, leave=True, initial=0):
         if self.pg_bar is None:
-            self.pg_bar = tqdm( total=total, desc=desc, leave=leave, ascii=True )
+            self.pg_bar = tqdm( total=total, desc=desc, leave=leave, ascii=True, initial=initial )
         else: print("progress_bar: already set.")
 
     def progress_bar_inc(self, c):
@@ -154,8 +154,8 @@ class InteractBase(object):
             self.pg_bar = None
         else: print("progress_bar not set.")
 
-    def progress_bar_generator(self, data, desc, leave=True):
+    def progress_bar_generator(self, data, desc, leave=True, initial=0):
-        self.pg_bar = tqdm( data, desc=desc, leave=leave, ascii=True )
+        self.pg_bar = tqdm( data, desc=desc, leave=leave, ascii=True, initial=initial )
         for x in self.pg_bar:
             yield x
         self.pg_bar.close()

mainscripts/Converter.py

@@ -2,6 +2,7 @@
 import multiprocessing
 import operator
 import os
+import pickle
 import shutil
 import sys
 import time
@@ -13,7 +14,8 @@ import numpy as np
 import numpy.linalg as npla
 
 import imagelib
-from converters import FrameInfo, ConverterConfig, ConvertMasked, ConvertFaceAvatar
+from converters import (ConverterConfig, ConvertFaceAvatar, ConvertMasked,
+                        FrameInfo)
 from facelib import FaceType, FANSegmentator, LandmarksProcessor
 from interact import interact as io
 from joblib import SubprocessFunctionCaller, Subprocessor
@@ -177,7 +179,7 @@ class ConvertSubprocessor(Subprocessor):
             return pf.frame_info.filename
 
     #override
-    def __init__(self, is_interactive, converter_config, frames, output_path):
+    def __init__(self, is_interactive, converter_config, frames, output_path, model_iter):
         if len (frames) == 0:
             raise ValueError ("len (frames) == 0")
 
@@ -203,20 +205,84 @@ class ConvertSubprocessor(Subprocessor):
 
         self.dcscn_host, self.superres_func = SubprocessFunctionCaller.make_pair(superres_func)
 
-        self.frames = frames
         self.output_path = output_path
+        self.model_iter = model_iter
+
         self.prefetch_frame_count = self.process_count = min(6,multiprocessing.cpu_count())
 
+        session_data = None
+        session_dat_path = self.output_path / 'session.dat'
+        if session_dat_path.exists():
+
+            try:
+                with open( str(session_dat_path), "rb") as f:
+                    session_data = pickle.loads(f.read())
+            except Exception as e:
+                pass
+
+        self.frames = frames
         self.frames_idxs = [ *range(len(self.frames)) ]
         self.frames_done_idxs = []
 
-        digits = [ str(i) for i in range(10)]
+        if self.is_interactive and session_data is not None:
+            s_frames = session_data.get('frames', None)
+            s_frames_idxs = session_data.get('frames_idxs', None)
+            s_frames_done_idxs = session_data.get('frames_done_idxs', None)
+            s_model_iter = session_data.get('model_iter', None)
+
+            frames_equal = (s_frames is not None) and \
+                           (s_frames_idxs is not None) and \
+                           (s_frames_done_idxs is not None) and \
+                           (s_model_iter is not None) and \
+                           (len(frames) == len(s_frames))
+
+            if frames_equal:
+                for i in range(len(frames)):
+                    frame = frames[i]
+                    s_frame = s_frames[i]
+                    if frame.frame_info.filename != s_frame.frame_info.filename:
+                        frames_equal = False
+                    if not frames_equal:
+                        break
+
+            if frames_equal:
+                io.log_info ("Using saved session.")
+                self.frames = s_frames
+                self.frames_idxs = s_frames_idxs
+                self.frames_done_idxs = s_frames_done_idxs
+
+                if self.model_iter != s_model_iter:
+                    #model is more trained, recompute all frames
+                    for frame in self.frames:
+                        frame.is_done = False
+
+                if self.model_iter != s_model_iter or \
+                    len(self.frames_idxs) == 0:
+                    #rewind to begin if model is more trained or all frames are done
+                   
+                    while len(self.frames_done_idxs) > 0:
+                        prev_frame = self.frames[self.frames_done_idxs.pop()]
+                        self.frames_idxs.insert(0, prev_frame.idx)
+
+                if len(self.frames_idxs) != 0:
+                    cur_frame = self.frames[self.frames_idxs[0]]
+                    cur_frame.is_shown = False
+
+            if not frames_equal:
+                session_data = None
+
+        if session_data is None:
+            for filename in Path_utils.get_image_paths(self.output_path): #remove all images in output_path
+                Path(filename).unlink()
+
+            frames[0].cfg = self.converter_config.copy()
+
         for i in range( len(self.frames) ):
             frame = self.frames[i]
             frame.idx = i
             frame.output_filename = self.output_path / ( Path(frame.frame_info.filename).stem + '.png' )
 
-        frames[0].cfg = self.converter_config.copy()
+
 
     #override
     def process_info_generator(self):
@@ -232,7 +298,7 @@ class ConvertSubprocessor(Subprocessor):
 
     #overridable optional
     def on_clients_initialized(self):
-        io.progress_bar ("Converting", len (self.frames_idxs) )
+        io.progress_bar ("Converting", len (self.frames_idxs), initial=len(self.frames_done_idxs) )
 
         self.process_remain_frames = not self.is_interactive
         self.is_interactive_quitting = not self.is_interactive
@@ -251,11 +317,25 @@ class ConvertSubprocessor(Subprocessor):
 
     #overridable optional
     def on_clients_finalized(self):
+        io.progress_bar_close()
 
         if self.is_interactive:
             self.screen_manager.finalize()
 
-        io.progress_bar_close()
+            for frame in self.frames:
+                frame.output_filename = None
+                frame.image = None
+
+            session_data = {
+                'frames': self.frames,
+                'frames_idxs': self.frames_idxs,
+                'frames_done_idxs': self.frames_done_idxs,
+                'model_iter' : self.model_iter,
+            }
+            save_path = self.output_path / 'session.dat'
+            save_path.write_bytes( pickle.dumps(session_data) )
+
+            io.log_info ("Session is saved to " + '/'.join (save_path.parts[-2:]) )
 
     cfg_change_keys = ['`','1', '2', '3', '4', '5', '6', '7', '8', '9',
                                  'q', 'a', 'w', 's', 'e', 'd', 'r', 'f', 't', 'g','y','h','u','j',
@@ -501,10 +581,7 @@ def main (args, device_args):
             io.log_err('Input directory not found. Please ensure it exists.')
             return
 
-        if output_path.exists():
+        if not output_path.exists():
-            for filename in Path_utils.get_image_paths(output_path):
-                Path(filename).unlink()
-        else:
             output_path.mkdir(parents=True, exist_ok=True)
 
         if not model_path.exists():
@@ -643,6 +720,7 @@ def main (args, device_args):
                         converter_config       = cfg,
                         frames                 = frames,
                         output_path            = output_path,
+                        model_iter             = model.get_iter()
                     ).run()
 
         model.finalize()