From 71bcd863e1d9b51fbce920de95816268119af9fb Mon Sep 17 00:00:00 2001 From: torzdf <36920800+torzdf@users.noreply.github.com> Date: Tue, 11 Jun 2019 17:26:43 +0000 Subject: [PATCH] Bugfixes and minor changes Fix: Occasional memory leak in convert Fix: Dynamic config loading in Windows Change: Better folder naming for snapshots --- lib/utils.py | 18 ++++++++++++++++++ plugins/convert/_config.py | 3 ++- plugins/extract/_config.py | 3 ++- plugins/train/_config.py | 3 ++- plugins/train/model/_base.py | 2 +- scripts/convert.py | 14 ++++++++++++-- 6 files changed, 37 insertions(+), 6 deletions(-) diff --git a/lib/utils.py b/lib/utils.py index f30fe39935..34c150b808 100644 --- a/lib/utils.py +++ b/lib/utils.py @@ -65,6 +65,24 @@ def get_image_paths(directory): return dir_contents +def full_path_split(path): + """ Split a given path into all of it's separate components """ + allparts = list() + while True: + parts = os.path.split(path) + if parts[0] == path: # sentinel for absolute paths + allparts.insert(0, parts[0]) + break + elif parts[1] == path: # sentinel for relative paths + allparts.insert(0, parts[1]) + break + else: + path = parts[0] + allparts.insert(0, parts[1]) + logger.trace("path: %s, allparts: %s", path, allparts) + return allparts + + def cv2_read_img(filename, raise_error=False): """ Read an image with cv2 and check that an image was actually loaded. Logs an error if the image returned is None. or an error has occured. diff --git a/plugins/convert/_config.py b/plugins/convert/_config.py index 88b9a9aac3..5e916e699d 100644 --- a/plugins/convert/_config.py +++ b/plugins/convert/_config.py @@ -8,6 +8,7 @@ from importlib import import_module from lib.config import FaceswapConfig +from lib.utils import full_path_split logger = logging.getLogger(__name__) # pylint: disable=invalid-name @@ -24,7 +25,7 @@ def set_defaults(self): if not default_files: continue base_path = os.path.dirname(os.path.realpath(sys.argv[0])) - import_path = dirpath.replace(base_path, "").replace("/", ".")[1:] + import_path = ".".join(full_path_split(dirpath.replace(base_path, ""))[1:]) plugin_type = import_path.split(".")[-1] for filename in default_files: self.load_module(filename, import_path, plugin_type) diff --git a/plugins/extract/_config.py b/plugins/extract/_config.py index 5cfdc8a4a7..2cb5e3fbdd 100644 --- a/plugins/extract/_config.py +++ b/plugins/extract/_config.py @@ -7,6 +7,7 @@ from importlib import import_module from lib.config import FaceswapConfig +from lib.utils import full_path_split logger = logging.getLogger(__name__) # pylint: disable=invalid-name @@ -23,7 +24,7 @@ def set_defaults(self): if not default_files: continue base_path = os.path.dirname(os.path.realpath(sys.argv[0])) - import_path = dirpath.replace(base_path, "").replace("/", ".")[1:] + import_path = ".".join(full_path_split(dirpath.replace(base_path, ""))[1:]) plugin_type = import_path.split(".")[-1] for filename in default_files: self.load_module(filename, import_path, plugin_type) diff --git a/plugins/train/_config.py b/plugins/train/_config.py index f92c07cb4b..ec67a99267 100644 --- a/plugins/train/_config.py +++ b/plugins/train/_config.py @@ -9,6 +9,7 @@ from lib.config import FaceswapConfig from lib.model.masks import get_available_masks +from lib.utils import full_path_split logger = logging.getLogger(__name__) # pylint: disable=invalid-name @@ -29,7 +30,7 @@ def set_defaults(self): if not default_files: continue base_path = os.path.dirname(os.path.realpath(sys.argv[0])) - import_path = dirpath.replace(base_path, "").replace("/", ".")[1:] + import_path = ".".join(full_path_split(dirpath.replace(base_path, ""))[1:]) plugin_type = import_path.split(".")[-1] for filename in default_files: self.load_module(filename, import_path, plugin_type) diff --git a/plugins/train/model/_base.py b/plugins/train/model/_base.py index 04a2f05c93..4f20462447 100644 --- a/plugins/train/model/_base.py +++ b/plugins/train/model/_base.py @@ -426,7 +426,7 @@ def snapshot_models(self): """ Take a snapshot of the model at current state and back up """ logger.info("Saving snapshot") src = self.model_dir - dst = get_folder("{}_{}".format(self.model_dir, self.iterations)) + dst = get_folder("{}_snapshot_{}_iters".format(self.model_dir, self.iterations)) for filename in os.listdir(src): if filename.endswith(".bk"): continue diff --git a/scripts/convert.py b/scripts/convert.py index 9dc7911171..fc598dd04f 100644 --- a/scripts/convert.py +++ b/scripts/convert.py @@ -492,12 +492,19 @@ def get_trainer(self, model_dir): def predict_faces(self): """ Get detected faces from images """ faces_seen = 0 + consecutive_no_faces = 0 batch = list() while True: item = self.in_queue.get() if item != "EOF": logger.trace("Got from queue: '%s'", item["filename"]) faces_count = len(item["detected_faces"]) + + # Safety measure. If a large stream of frames appear that do not have faces, + # these will stack up into RAM. Keep a count of consecutive frames with no faces. + # If self.batchsize number of frames appear, force the current batch through + # to clear RAM. + consecutive_no_faces = consecutive_no_faces + 1 if faces_count == 0 else 0 self.faces_count += faces_count if faces_count > 1: self.verify_output = True @@ -509,8 +516,10 @@ def predict_faces(self): faces_seen += faces_count batch.append(item) - if faces_seen < self.batchsize and item != "EOF": - logger.trace("Continuing. Current batchsize: %s", faces_seen) + if item != "EOF" and (faces_seen < self.batchsize and + consecutive_no_faces < self.batchsize): + logger.trace("Continuing. Current batchsize: %s, consecutive_no_faces: %s", + faces_seen, consecutive_no_faces) continue if batch: @@ -526,6 +535,7 @@ def predict_faces(self): self.queue_out_frames(batch, predicted) + consecutive_no_faces = 0 faces_seen = 0 batch = list() if item == "EOF":