added a bunch of stuff, release on the way

.gitignore

@@ -2,5 +2,4 @@ venv
 main.build
 dist
 test
-1x-WB-Denoise.pth
 test.py

README.md

@@ -1,11 +1,3 @@
-     _____       __  __       
-    |  __ \     |  \/  |      
-    | |__) |   _| \  / | ___  
-    |  ___/ | | | |\/| |/ _ \ 
-    | |   | |_| | |  | | (_) |
-    |_|    \__, |_|  |_|\___/ 
-            __/ |             
-           |___/
 # PyMo (Python Motion Visualizer GUI)
 ### A CLI script for visualizing differences between video frames, eg. motion.
 
@@ -41,24 +33,24 @@
 
 ## Requirements:
 - Python 3.10 or newer
-- Nvidia GPU with installed drivers (Optional but Recommended)
+- Nvidia GPU
+  - Minimum: GTX 1650 4GB
+  - Recommended: RTX 4060 8GB
 - At least 16GB of RAM
 
 ---
 
 ## Installation:
 ```bash
-git clone https://github.com/TheElevatedOne/pymo.git && cd pymo
-python3 -m venv venv
-source venv/bin/activate
-pip install -r requirements.txt
+git clone -b 1.0.0 https://github.com/TheElevatedOne/pymo.git && cd pymo
+./install-venv.sh
 ```
 
 ---
 
 ## Running:
 ```
-usage: main.py [-h] -i INPUT [-o OUTPUT] [-n NAME] [-f OFFSET] [-t THREADS] [-s] [-c] [-m MODEL]
+usage: pymo [-h] -i INPUT [-o PATH] [-n NAME] [-f INT[1, 50]] [-t CPU[1, 8]] [-s] [-c] [-m {knn,nlm,sr}] [-sr {Directory Empty!}]
 
      _____       __  __       
     |  __ \     |  \/  |      
@@ -68,30 +60,47 @@ usage: main.py [-h] -i INPUT [-o OUTPUT] [-n NAME] [-f OFFSET] [-t THREADS] [-s]
     |_|    \__, |_|  |_|\___/ 
             __/ |             
            |___/              
-           
+
    Python Motion Visualizer CLI
 
 options:
   -h, --help            show this help message and exit
+
+Required:
+  Required Arguments
+
   -i INPUT, --input INPUT
                         Relative path to the Input Video
-  -o OUTPUT, --output OUTPUT
+
+Optional:
+  Optional Arguments
+
+  -o PATH, --output PATH
                         (Optional) Absolute path to output directory
   -n NAME, --name NAME  (Optional) Custom Filename for the video
-  -f OFFSET, --offset OFFSET
+  -f INT[1, 50], --offset INT[1, 50]
                         (Optional) Number of Offset Frames [Default = 5]
-  -t THREADS, --threads THREADS
+  -t CPU[1, 8], --threads CPU[1, 8]
                         (Optional) Amount of threads to run the process on [Default = 2]
   -s, --slow_motion     (Optional) Sets the FPS of the Output Video to half the original;
-                        Essentially creating a slow-motion of the original without interpolation
-  -c, --cpu             (Optional) Denoising step by default runs on CUDA Acceleration (if Nvidia GPU Available)
-                        Setting this makes it run on CPU even if GPU is Available
-  -m MODEL, --model MODEL
-                        (Optional) Model to use when denoising via GPU [Default = knn] [Options: nlm, knn]
+                                   Essentially creating a slow-motion of the original without interpolation
+
+Denoising:
+  Denoising Arguments
+
+  -c, --cpu             (Optional) Denoising step by default runs on CUDA Acceleration (if Nvidia GPU Available);
+                                   Setting this makes it run on CPU even if GPU is Available
+  -m {knn,nlm,sr}, --model {knn,nlm,sr}
+                        (Optional) Model to use when denoising via GPU [Default = knn];
+                                   SR - Super Resolution (ESRGAN, SwinIR, ...), 
+                                   Needs a PyTorch weight file (.pth) in ./weights/ to be active.
+  -sr {Directory Empty!}, --super_resolution {Directory Empty!}
+                        (Optional) Choosing Weights for SR (if available)
+                                   [Default = Directory Empty!]
+
 ```
 
-### Detailed info on arguments:
-- `-i/--input [FILE]` A Video File with OpenCV2 supported extensions
+### Detailed info on some arguments:
 - `-o/--output [PATH]` An absolute path to a directory, e.g. on Linux `/home/${whoami}/directory`
   - If left empty, it will default to Input Video Dir
 - `-n/--name [NAME]` Custom name for the output video, without extension. Default is `pymo_{input-vid-name}`
@@ -108,7 +117,33 @@ options:
   - ex. 60 fps Video -> 30 fps Video but Twice the length
 - `-c/--cpu` The program uses [PyCuda](https://pypi.org/project/pycuda/) for GPU Acceleration on Denoising by default (unless it does not detect a GPU)
   - This forces the Program to use the CPU, which results in higher quality render, but with a long render time.
-- `-m/--model` GPU Denoising Model switcher, both are good for some things
+- `-m/--model` GPU Denoising Model switcher.
+  - KNN/NLM2 are good for some things, they are fast but sacrifice quality.
+  - SR (Super Resolution) models (or weights) are slower than the above, but faster than CPU and have amazing quality.
+    - They are trained weights on datasets (AI), which can do many things. I mainly use them to Denoise.
+    - They need fast GPU's to run and a lot of VRAM. This may be the choice for someone with a good GPU.
+    - Weights are not packaged with the program, Download them [here](./weights/README.md).
+
+---
+
+## Optimization:
+As I develop this project, I try to optimize any algorithm I use after I got it working.<br>
+With that, I just make 80% improvement on the Motion Generation (Difference Blend) and 50% improvement on CPU denoising.
+
+Current Times (On My Machine):
+- Specs:
+  - CPU - Intel i5-9300H
+  - GPU - Nvidia GTX 1650 Mobile
+- Times (for Compute Intesive Tasks) [using test_01.mp4 | 4K, 210 Frames]:
+  - Exporting Frames of the Video - 44s
+  - Motion Generation - 1m 11s
+  - Contrast Filter - 27s
+  - Denoising:
+    - CPU - 6m 2s
+    - KNN - 44s
+    - NLM2 - 54s
+    - SR (1x-WB-Denoise.pth) - 6m 52s
+  - Encoding Video - 17s
 
 ---
 
@@ -120,12 +155,15 @@ options:
 ## Modules and scripts used:
 - [**tqdm**](https://pypi.org/project/tqdm/) for progress bars
 - [**PIL/Pillow**](https://pypi.org/project/pillow/) for easy image transformation
-- [**blend-modes**](https://pypi.org/project/blend-modes/) for the difference blend
 - [**opencv-python**](https://pypi.org/project/opencv-python/) for simpler video transformation, reading, writing and denoising
-- [**numpy**](https://pypi.org/project/numpy/) for converting Pillow images to cv2 arrays and vice-versa
-- [**pycuda**](https://pypi.org/project/pycuda/) for running GPU-accelerated denoising
-- [**torch**](https://pypi.org/project/torch/) just for checking of user has GPU or not
+- [**imageio**](https://pypi.org/project/imageio/) for reading and writing images
+- [**numpy**](https://pypi.org/project/numpy/) for converting Pillow images to cv2 arrays and vice versa
+- [**pycuda**](https://pypi.org/project/pycuda/) for running GPU-accelerated Denoising
+- [**torch**](https://pypi.org/project/torch/) for running GPU Super Resolution Models for Denoising
+- [**torchvision**](https://pypi.org/project/torchvision/) for converting numpy arrays into tensors and vice versa
+- [**spandrel**](https://pypi.org/project/spandrel/) project of chaiNNer developer, for reading all types of SR weights
 - [**PyCuda_Denoise_Filters**](https://github.com/AlainPaillou/PyCuda_Denoise_Filters) for GPU denoising
   - Both of the `Mono` scripts were used in my project and were modified to fit the use-case
+- [**Nuitka**](https://pypi.org/project/Nuitka/) for compiling `main.py` into a binary for easier running and adding to `$PATH`
 <br>
 <br>

main.py

@@ -1,37 +1,69 @@
 import argparse
 import multiprocessing
+import os.path
 import sys
+import termcolor
+import traceback
 from src.differ import Difference
 
 
 def parse() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="""
-     _____       __  __       
-    |  __ \     |  \/  |      
-    | |__) |   _| \  / | ___  
-    |  ___/ | | | |\/| |/ _ \ 
-    | |   | |_| | |  | | (_) |
-    |_|    \__, |_|  |_|\___/ 
-            __/ |             
-           |___/              
-           
-   Python Motion Visualizer CLI""", formatter_class=argparse.RawTextHelpFormatter)
+    """Function for parsing CLI arguments"""
+    sr_models = [
+        os.path.splitext(os.path.basename(x))[0] for x in os.listdir(
+            os.path.join(
+                os.path.dirname(__file__),
+                "weights"
+            )
+        ) if ".pth" in x
+    ]
+    sr_models.sort()
 
-    parser.add_argument("-i", "--input", required=True, help="""Relative path to the Input Video""", type=str)
-    parser.add_argument("-o", "--output", required=False, help="(Optional) Absolute path to output directory", type=str,
-                        metavar="PATH")
-    parser.add_argument("-n", "--name", required=False, help="(Optional) Custom Filename for the video", type=str)
-    parser.add_argument("-f", "--offset", required=False, default=5, help="(Optional) Number of Offset Frames [Default = 5]",
-                        type=int, choices=range(1, 50), metavar="INT[1, 50]")
-    parser.add_argument("-t", "--threads", required=False, default=2, help="(Optional) Amount of threads to run the process on [Default = 2]",
-                        type=int, choices=range(1, multiprocessing.cpu_count()), metavar="CPU[1, %d]" % multiprocessing.cpu_count())
-    parser.add_argument("-m", "--model", required=False, default="knn",
-                        help="(Optional) Model to use when denoising via GPU [Default = knn]",
-                        choices=["knn", "nlm"])
-    parser.add_argument("-s", "--slow_motion", required=False, default=False, help="""(Optional) Sets the FPS of the Output Video to half the original;
+    if not len(sr_models):
+        sr_models.append("Directory Empty!")
+
+    parser = argparse.ArgumentParser(description="    " + termcolor.colored(" _____       ", "light_blue", attrs=["bold"]) + termcolor.colored("__  __       ", "light_green", attrs=["bold"]) + "\n" + \
+    "    " + termcolor.colored("|  __ \     ", "light_blue", attrs=["bold"]) + termcolor.colored("|  \/  |      ", "light_green", attrs=["bold"]) + "\n" + \
+    "    " + termcolor.colored("| |__) |   _", "light_blue", attrs=["bold"]) + termcolor.colored("| \  / | ___  ", "light_green", attrs=["bold"]) + "\n" + \
+    "    " + termcolor.colored("|  ___/ | | |", "light_blue", attrs=["bold"]) + termcolor.colored(" |\/| |/ _ \ ", "light_green", attrs=["bold"]) + "\n" + \
+    "    " + termcolor.colored("| |   | |_| |", "light_blue", attrs=["bold"]) + termcolor.colored(" |  | | (_) |", "light_green", attrs=["bold"]) + "\n" + \
+    "    " + termcolor.colored("|_|    \__, |", "light_blue", attrs=["bold"]) + termcolor.colored("_|  |_|\___/ ", "light_green", attrs=["bold"]) + "\n" + \
+    "    " + termcolor.colored("        __/ |", "light_blue", attrs=["bold"]) + termcolor.colored("             ", "light_green", attrs=["bold"]) + "\n" + \
+    "    " + termcolor.colored("       |___/ ", "light_blue", attrs=["bold"]) + termcolor.colored("             ", "light_green", attrs=["bold"]) + "\n\n" + \
+    "   " + termcolor.colored("Python ", "light_blue", attrs=["bold"])  + termcolor.colored("Motion ", "light_green", attrs=["bold"]) + \
+    termcolor.colored("Visualizer CLI", "white", attrs=["bold"]), formatter_class=argparse.RawTextHelpFormatter)
+
+    required = parser.add_argument_group("Required", "Required Arguments")
+    optional = parser.add_argument_group("Optional", "Optional Arguments")
+    denoise = parser.add_argument_group("Denoising", description="Denoising Arguments")
+    required.add_argument("-i", "--input", required=True,
+                        help="""Relative path to the Input Video""", type=str)
+    optional.add_argument("-o", "--output", required=False,
+                        help="(Optional) Absolute path to output directory", type=str,metavar="PATH")
+    optional.add_argument("-n", "--name", required=False, help="(Optional) Custom Filename for the video",
+                        type=str)
+    optional.add_argument("-f", "--offset", required=False, default=5,
+                        help="(Optional) Number of Offset Frames [Default = 5]",type=int, choices=range(1, 50),
+                        metavar="INT[1, 50]")
+    optional.add_argument("-t", "--threads", required=False, default=2,
+                        help="(Optional) Amount of threads to run the process on [Default = 2]",
+                        type=int, choices=range(1, multiprocessing.cpu_count()),
+                        metavar="CPU[1, %d]" % multiprocessing.cpu_count())
+    optional.add_argument("-s", "--slow_motion", required=False, default=False,
+                        help="""(Optional) Sets the FPS of the Output Video to half the original;
            Essentially creating a slow-motion of the original without interpolation""", action="store_true")
-    parser.add_argument("-c", "--cpu", required=False, default=False, help="""(Optional) Denoising step by default runs on CUDA Acceleration (if Nvidia GPU Available);
+    denoise.add_argument("-c", "--cpu", required=False, default=False,
+                         help="""(Optional) Denoising step by default runs on CUDA Acceleration (if Nvidia GPU Available);
            Setting this makes it run on CPU even if GPU is Available""", action="store_true")
+    denoise.add_argument("-m", "--model", required=False, default="knn",
+                        help="""(Optional) Model to use when denoising via GPU [Default = knn];
+           SR - Super Resolution (ESRGAN, SwinIR, ...), 
+           Needs a PyTorch weight file (.pth) in ./weights/ to be active.""",
+                        choices=["knn", "nlm", "sr"])
+    denoise.add_argument("-sr", "--super_resolution", required=False,
+                        help=f'''(Optional) Choosing Weights for SR (if available)
+           [Default = {sr_models[0]}]''',
+                        choices=sr_models, default=sr_models[0])
 
     if len(sys.argv) == 1:
         parser.print_help(sys.stderr)
@@ -41,18 +73,19 @@ def parse() -> argparse.Namespace:
 
 def main() -> None:
     args = parse() # Parse CLI Arguments
-    print("""
-     _____       __  __       
-    |  __ \     |  \/  |      
-    | |__) |   _| \  / | ___  
-    |  ___/ | | | |\/| |/ _ \ 
-    | |   | |_| | |  | | (_) |
-    |_|    \__, |_|  |_|\___/ 
-            __/ |             
-           |___/              
-----------------------------------""")
+    print("    " + termcolor.colored(" _____       ", "light_blue", attrs=["bold"]) + termcolor.colored("__  __       ", "light_green", attrs=["bold"]) + "\n" + \
+    "    " + termcolor.colored("|  __ \     ", "light_blue", attrs=["bold"]) + termcolor.colored("|  \/  |      ", "light_green", attrs=["bold"]) + "\n" + \
+    "    " + termcolor.colored("| |__) |   _", "light_blue", attrs=["bold"]) + termcolor.colored("| \  / | ___  ", "light_green", attrs=["bold"]) + "\n" + \
+    "    " + termcolor.colored("|  ___/ | | |", "light_blue", attrs=["bold"]) + termcolor.colored(" |\/| |/ _ \ ", "light_green", attrs=["bold"]) + "\n" + \
+    "    " + termcolor.colored("| |   | |_| |", "light_blue", attrs=["bold"]) + termcolor.colored(" |  | | (_) |", "light_green", attrs=["bold"]) + "\n" + \
+    "    " + termcolor.colored("|_|    \__, |", "light_blue", attrs=["bold"]) + termcolor.colored("_|  |_|\___/ ", "light_green", attrs=["bold"]) + "\n" + \
+    "    " + termcolor.colored("        __/ |", "light_blue", attrs=["bold"]) + termcolor.colored("             ", "light_green", attrs=["bold"]) + "\n" + \
+    "    " + termcolor.colored("       |___/ ", "light_blue", attrs=["bold"]) + termcolor.colored("             ", "light_green", attrs=["bold"]) + "\n\n" + \
+    "   " + termcolor.colored("Python ", "light_blue", attrs=["bold"])  + termcolor.colored("Motion ", "light_green", attrs=["bold"]) + \
+    termcolor.colored("Visualizer CLI", "white", attrs=["bold"]) + "\n")
     print("Initializing")
-    di = Difference(args.input, args.output, args.name, args.offset, args.threads, args.slow_motion, args.cpu, args.model) # Initialize imported module
+    di = Difference(args.input, args.output, args.name, args.offset, args.threads,
+                    args.slow_motion, args.cpu, args.model, args.super_resolution) # Initialize imported module
     print("Creating Temporary Files Directory")
     print("----------------------------------")
     try:
@@ -62,15 +95,22 @@ def main() -> None:
         print("02. Difference Generation")
         di.dif_frames()
         print("03. Post Processing")
-        di.gen_video()
+        output = di.gen_video()
         print("    02. Removing Temporary Directory")
         di.temp_dir()
+        print("----------------------------------")
+        print(termcolor.colored("Finished!", "light_green", attrs=["bold"]))
+        print(f"Video rendered at {termcolor.colored(output, color='light_yellow')}")
     except Exception as e:
         di.temp_dir()
         print("----------------------------------")
-        print("Program Failed on Runtime due to Exception:")
-        print(e)
+        print(f">>> {termcolor.colored('Program Failed on Runtime due to Exception:', 'red', 'on_black', ['bold'])}")
+        print(f">>> {termcolor.colored(e, on_color='on_black', attrs=['bold'])}")
+        traceback.print_exc()
     except KeyboardInterrupt:
+        print("----------------------------------")
+        print(f">>> {termcolor.colored('KeyboardInterrupt', 'cyan', attrs=['bold'])}")
+        print(">>> Removing Temporary Directory")
         di.temp_dir()
         sys.exit(130)
 

requirements.txt

@@ -5,4 +5,6 @@ pycuda>=2024.1.2
 torch>=2.5.1
 tqdm>=4.67.1
 imageio>=2.36.1
-nuitka>=2.5.8
+torchvision>=0.8.0
+spandrel>=0.4.0
+nuitka>=2.5.8

src/cuda/pycuda_denoise.py

@@ -9,7 +9,7 @@ class PyCudaDenoise:
     Based on KNN filter CUDA program provided by Nvidia - CUDA SDK samples
     Alain PAILLOU - August 2019
     TheElevatedOne - December 2024 - Modified for this CLI Program"""
-    def __init__(self, img: cv2.Mat, model: str):
+    def __init__(self, img: cv2.Mat, model: str) -> None:
         drv.init()
         device = drv.Device(0)
         self.ctx = device.make_context()
@@ -200,7 +200,7 @@ def __init__(self, img: cv2.Mat, model: str):
 """)
         self.model = model
 
-    def run(self):
+    def run(self) -> np.ndarray:
         if self.model == "nlm":
             mono_gpu = self.nlm2.get_function("NLM2_Mono")
         else: