From d0b32b56faa7bf0e29637120ce7459f20dbc5229 Mon Sep 17 00:00:00 2001 From: n00mkrad Date: Sun, 22 Aug 2021 17:27:39 +0200 Subject: [PATCH] actually fixed model comparison tool, more fixes --- Code/Forms/ModelComparisonForm.cs | 20 ++++++++++++------ Code/Forms/ModelSelectForm.cs | 31 ++++++++++++++++++++-------- Code/IO/IOUtils.cs | 27 ++++++++++++++++++++---- Code/IO/Paths.cs | 2 ++ Code/Main/MainForm.cs | 19 +++++++++-------- Code/Main/PostProcessing.cs | 10 ++++----- Code/OS/NcnnUtils.cs | 34 +++++++++++++++++++++++++++++++ Code/UI/PreviewUI.cs | 22 ++++++++++++++++---- 8 files changed, 128 insertions(+), 37 deletions(-) diff --git a/Code/Forms/ModelComparisonForm.cs b/Code/Forms/ModelComparisonForm.cs index 558c996..4c0a1bf 100644 --- a/Code/Forms/ModelComparisonForm.cs +++ b/Code/Forms/ModelComparisonForm.cs @@ -68,6 +68,7 @@ private async void runBtn_Click(object sender, EventArgs e) Program.ShowMessage("No image loaded!", "Error"); return; } + Enabled = false; cutoutMode = cropMode.SelectedIndex == 1; if (cutoutMode) @@ -80,20 +81,28 @@ private async void runBtn_Click(object sender, EventArgs e) { currentSourcePath = Paths.tempImgPath; } + string[] lines = Regex.Split(modelPathsBox.Text, "\r\n|\r|\n"); + if(comparisonMode.SelectedIndex == 0) { string outpath = Path.Combine(Paths.imgOutPath, "!Original.png"); await ImageProcessing.ConvertImage(currentSourcePath, GetSaveFormat(), false, ImageProcessing.ExtMode.UseNew, false, outpath); await ProcessImage(outpath, "Original"); } + for (int i = 0; i < lines.Length; i++) { - if (!File.Exists(lines[i])) + if (!IoUtils.IsPathValid(lines[i])) + { + Logger.Log($"ModelComparisonForm: Path '{lines[i]}' is not valid, skipping upscale with this model!"); continue; + } + ModelData mdl = new ModelData(lines[i], null, ModelData.ModelMode.Single); await DoUpscale(i, mdl, !cutoutMode); } + bool vert = compositionMode.SelectedIndex == 1; MagickImage merged = ImgUtils.MergeImages(Directory.GetFiles(Paths.imgOutPath, "*.png", SearchOption.AllDirectories), vert, true); string mergedPath = Path.Combine(Paths.imgOutPath, Path.GetFileNameWithoutExtension(Program.lastImgPath) + "-composition"); @@ -108,10 +117,13 @@ private async void runBtn_Click(object sender, EventArgs e) static ImageProcessing.Format GetSaveFormat() { ImageProcessing.Format saveFormat = ImageProcessing.Format.PngFast; + if (Config.GetInt("previewFormat") == 1) saveFormat = ImageProcessing.Format.Jpeg; + if (Config.GetInt("previewFormat") == 2) saveFormat = ImageProcessing.Format.Weppy; + return saveFormat; } @@ -148,17 +160,13 @@ async Task DoUpscale (int index, ModelData mdl, bool fullImage) if (fullImage) inpath = Paths.tempImgPath.GetParentDir(); await Upscale.Run(inpath, Paths.compositionOut, mdl, false, Config.GetBool("alpha"), PreviewUi.PreviewMode.None); outImg = Directory.GetFiles(Paths.compositionOut, "*.png", SearchOption.AllDirectories)[0]; - await PostProcessing.PostprocessingSingle(outImg, false); + await PostProcessing.PostprocessingSingle(outImg, false, 10, false); await ProcessImage(PreviewUi.lastOutfile, mdl.model1Name); IoUtils.TryCopy(PreviewUi.lastOutfile, Path.Combine(Paths.imgOutPath, $"{index}-{mdl.model1Name}.png"), true); } catch (Exception e) { if (Program.canceled) return; - - if (e.Message.ToLower().Contains("index")) - Program.ShowMessage("The upscale process seems to have exited before completion!", "Error"); - Program.Cancel($"An error occured: {e.Message}"); } diff --git a/Code/Forms/ModelSelectForm.cs b/Code/Forms/ModelSelectForm.cs index e9971a4..c913246 100644 --- a/Code/Forms/ModelSelectForm.cs +++ b/Code/Forms/ModelSelectForm.cs @@ -1,4 +1,5 @@ using Cupscale.IO; +using Cupscale.OS; using System; using System.Collections.Generic; using System.ComponentModel; @@ -85,25 +86,37 @@ private void ForceLowercaseExtensions(string path) { foreach (FileInfo file in IoUtils.GetFileInfosSorted(path, true, "*.*")) { - if(file.Extension == ".PTH") + if (file.Extension == ".PTH") file.MoveTo(Path.ChangeExtension(file.FullName, "pth")); } } - private void BuildTree(DirectoryInfo directoryInfo, TreeNodeCollection addInMe) + private void BuildTree(DirectoryInfo directoryInfo, TreeNodeCollection nodeCollection) { - TreeNode curNode = addInMe.Add(directoryInfo.Name); + TreeNode currNode = nodeCollection.Add(directoryInfo.Name); foreach (FileInfo file in directoryInfo.GetFiles()) { if (file.Extension == ".pth") // Hide any other file extension - curNode.Nodes.Add(file.FullName, Path.ChangeExtension(file.Name, null)); + currNode.Nodes.Add(file.FullName, Path.ChangeExtension(file.Name, null)); } - foreach (DirectoryInfo subdir in directoryInfo.GetDirectories()) + foreach (DirectoryInfo subDir in directoryInfo.GetDirectories()) { - if (subdir.GetFiles("*.pth", SearchOption.AllDirectories).Length > 0) // Don't list folders that have no PTH files - BuildTree(subdir, curNode.Nodes); + bool isNcnnModel = NcnnUtils.IsDirNcnnModel(subDir.FullName); + + if (isNcnnModel) + currNode.Nodes.Add(subDir.FullName, subDir.Name.Substring(0, subDir.Name.Length - 5)); + + bool hasAnyPthFiles = subDir.GetFiles("*.pth", SearchOption.AllDirectories).Length > 0; + bool hasAnyBinFiles = subDir.GetFiles("*.bin", SearchOption.AllDirectories).Length > 0; + bool hasAnyParamFiles = subDir.GetFiles("*.param", SearchOption.AllDirectories).Length > 0; + + if (isNcnnModel || subDir.Name.StartsWith(".")) + continue; // Don't add this folder to the tree if it's a model, not a dir with more models + + if (hasAnyPthFiles || hasAnyBinFiles || hasAnyParamFiles) // Don't list folders that have no model files + BuildTree(subDir, currNode.Nodes); } } @@ -132,7 +145,7 @@ private void cancelBtn_Click(object sender, EventArgs e) private void modelTree_AfterSelect(object sender, TreeViewEventArgs e) { - confirmBtn.Enabled = (Path.GetExtension(modelTree.SelectedNode.Name) == ".pth"); + confirmBtn.Enabled = (Path.GetExtension(modelTree.SelectedNode.Name) == ".pth" || modelTree.SelectedNode.Name.EndsWith(".ncnn")); } private void ModelSelectForm_KeyPress(object sender, KeyPressEventArgs e) @@ -140,7 +153,7 @@ private void ModelSelectForm_KeyPress(object sender, KeyPressEventArgs e) if (e.KeyChar == (char)Keys.Enter) { e.Handled = true; - if(confirmBtn.Enabled) + if (confirmBtn.Enabled) confirmBtn_Click(null, null); } diff --git a/Code/IO/IOUtils.cs b/Code/IO/IOUtils.cs index 3ead1f5..58b9267 100644 --- a/Code/IO/IOUtils.cs +++ b/Code/IO/IOUtils.cs @@ -77,17 +77,36 @@ public static bool IsPathDirectory(string path) public static bool IsFileValid(string path) { if (path == null) - { return false; - } + if (!File.Exists(path)) - { return false; - } return true; } + public static bool IsDirValid(string path) + { + if (path == null) + return false; + + if (!Directory.Exists(path)) + return false; + + return true; + } + + public static bool IsPathValid(string path) + { + if (path == null) + return false; + + if (IsPathDirectory(path)) + return IsDirValid(path); + else + return IsFileValid(path); + } + public static async Task CopyDir(string sourceDir, string targetDir, string wildcard = "*", bool move = false, bool onlyCompatibles = false, string removeFromName = "") { Logger.Log("[IOUtils] Copying directory \"" + sourceDir + "\" to \"" + targetDir + "\" (Move: " + move + " - RemoveFromName: " + removeFromName + ")"); diff --git a/Code/IO/Paths.cs b/Code/IO/Paths.cs index afbc6aa..9c6bdf2 100644 --- a/Code/IO/Paths.cs +++ b/Code/IO/Paths.cs @@ -23,6 +23,8 @@ internal class Paths public static readonly string pythonTuringUrl = "https://dl.nmkd-hz.de/flowframes/setupfiles/py-tu/v1/py-tu.7z"; public static readonly string pythonAmpereUrl = "https://dl.nmkd-hz.de/flowframes/setupfiles/py-amp/v1/py-amp.7z"; + public static readonly string ncnnMdlDir = ".ncnn-models"; + public static void Init() { binPath = Path.Combine(GetDataPath(), "bin"); diff --git a/Code/Main/MainForm.cs b/Code/Main/MainForm.cs index 3569283..8a6a946 100644 --- a/Code/Main/MainForm.cs +++ b/Code/Main/MainForm.cs @@ -122,7 +122,7 @@ void LoadLastUsedModels () { string mdl1 = Config.Get("lastMdl1"); - if (File.Exists(mdl1)) + if (File.Exists(mdl1) || (mdl1.EndsWith(".ncnn") && Directory.Exists(mdl1))) { Program.currentModel1 = mdl1; model1TreeBtn.Text = Path.GetFileNameWithoutExtension(Program.currentModel1); @@ -130,8 +130,8 @@ void LoadLastUsedModels () string mdl2 = Config.Get("lastMdl2"); - if (File.Exists(mdl2)) - { + if (File.Exists(mdl2) || (mdl2.EndsWith(".ncnn") && Directory.Exists(mdl2))) + { Program.currentModel2 = mdl2; model2TreeBtn.Text = Path.GetFileNameWithoutExtension(Program.currentModel2); } @@ -145,6 +145,7 @@ public async Task CheckInstallation () float lastProg; string lastStatus = ""; + public void SetProgress(float prog, string statusText = "") { if(prog >= 0 && lastProg != prog) @@ -433,15 +434,15 @@ private async void upscaleBtn_Click(object sender, EventArgs e) if (Program.busy) return; - if(Upscale.currentAi == Implementations.Imps.esrganPytorch && !PreviewUi.HasValidModelSelection()) - { - Program.ShowMessage("Invalid model selection.\nMake sure you have selected a model and that the file still exists.", "Error"); - return; - } + //if(Upscale.currentAi == Implementations.Imps.esrganPytorch && !PreviewUi.HasValidModelSelection()) + //{ + // Program.ShowMessage("Invalid model selection.\nMake sure you have selected a model and that the file still exists.", "Error"); + // return; + //} if (!PreviewUi.HasValidModelSelection()) { - Program.ShowMessage("Invalid model selection - NCNN does not support interpolation or chaining.", "Error"); + //Program.ShowMessage("Invalid model selection - NCNN does not support interpolation or chaining.", "Error"); return; } diff --git a/Code/Main/PostProcessing.cs b/Code/Main/PostProcessing.cs index c5d62cc..d1ff843 100644 --- a/Code/Main/PostProcessing.cs +++ b/Code/Main/PostProcessing.cs @@ -13,25 +13,25 @@ namespace Cupscale.Main { class PostProcessing { - public static async Task PostprocessingSingle(string path, bool dontResize = false, int retryCount = 20) + public static async Task PostprocessingSingle(string path, bool dontResize = false, int retryCount = 20, bool trimPngExt = true) { if (!IoUtils.IsFileValid(path)) return; string newPath = ""; - if (Path.GetExtension(path) != ".tmp") - newPath = path.Substring(0, path.Length - 8); - else + if (trimPngExt) newPath = path.Substring(0, path.Length - 4); + Logger.Log($"PostProc: Trimmed filename from '{Path.GetFileName(path)}' to '{Path.GetFileName(newPath)}'"); + try { File.Move(path, newPath); } catch (Exception e) // An I/O error can appear if the file is still locked by python (?) { - Logger.Log("Failed to move/rename! " + e.Message + "\n" + e.StackTrace); + Logger.Log($"Failed to move/rename! ('{path}' => '{newPath}') {e.Message}\n{e.StackTrace}"); if (retryCount > 0) { diff --git a/Code/OS/NcnnUtils.cs b/Code/OS/NcnnUtils.cs index b0fd324..a881b39 100644 --- a/Code/OS/NcnnUtils.cs +++ b/Code/OS/NcnnUtils.cs @@ -24,8 +24,16 @@ class NcnnUtils public static async Task ConvertNcnnModel(string modelPath, string filenamePattern) { + Logger.Log($"ConvertNcnnModel: {modelPath}"); + try { + if (IsDirNcnnModel(modelPath)) + { + currentNcnnModel = modelPath; + return; + } + string modelName = Path.GetFileName(modelPath); ncnnDir = Path.Combine(Config.Get("modelPath"), ".ncnn-models"); Directory.CreateDirectory(ncnnDir); @@ -109,6 +117,32 @@ private static void OutputHandler(object sendingProcess, DataReceivedEventArgs o lastNcnnOutput += $"{data}\n"; } + public static bool IsDirNcnnModel (string path, bool requireSuffix = true, bool noMoreThanTwoFiles = false) + { + try + { + if (!IoUtils.IsPathDirectory(path)) + return false; + + DirectoryInfo dir = new DirectoryInfo(path); + bool suffixValid = dir.Name.EndsWith(".ncnn"); + bool filesValid = dir.GetFiles("*.bin").Length == 1 && dir.GetFiles("*.param").Length == 1; + + if (noMoreThanTwoFiles && dir.GetFiles("*").Length > 2) + filesValid = false; + + if (requireSuffix && !suffixValid) + return false; + + return filesValid; + } + catch(Exception e) + { + Logger.Log($"IsDirNcnnModel Exception: {e.Message}. Defaulting to false."); + return false; + } + } + public static int GetNcnnModelScale(string modelDir) { try diff --git a/Code/UI/PreviewUI.cs b/Code/UI/PreviewUI.cs index ea7d469..3ac72d5 100644 --- a/Code/UI/PreviewUI.cs +++ b/Code/UI/PreviewUI.cs @@ -122,11 +122,21 @@ public static void TabSelected() public static bool HasValidModelSelection(bool showErrorMsgsIfInvalid = true) { Implementations.Implementation ai = Upscale.currentAi; + bool ncnn = ai == Implementations.Imps.esrganNcnn || ai == Implementations.Imps.realEsrganNcnn; - if(ai == Implementations.Imps.esrganPytorch) + if (ai == Implementations.Imps.esrganPytorch) { bool valid = true; + if (NcnnUtils.IsDirNcnnModel(Program.currentModel1) || NcnnUtils.IsDirNcnnModel(Program.currentModel2)) + valid = false; // NCNN models not compatible with pytorch + + if (!valid && showErrorMsgsIfInvalid) + { + Program.ShowMessage("Invalid model selection - You have selected one or more models that are not compatible with this implementation!", "Error"); + return false; + } + if (model1.Enabled && !File.Exists(Program.currentModel1)) valid = false; if (model2.Enabled && !File.Exists(Program.currentModel2)) @@ -138,14 +148,18 @@ public static bool HasValidModelSelection(bool showErrorMsgsIfInvalid = true) return valid; } - if (ai == Implementations.Imps.esrganNcnn) + if (ncnn) { bool valid = true; - valid = Program.mainForm.IsSingleModleMode(); + if (!Program.mainForm.IsSingleModleMode()) + valid = false; - if(!valid && showErrorMsgsIfInvalid) + if (!valid && showErrorMsgsIfInvalid) + { Program.ShowMessage("Invalid model selection - NCNN does not support interpolation or chaining.", "Error"); + return false; + } return valid; }