feat: clean code, add CI, precommit

.github/workflows/ci.yml

@@ -0,0 +1,18 @@
+name: CI
+
+on: [push, pull_request]
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    name: Do the code respects Python standards?
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: '3.10.0'
+      - name: Install pre-commit
+        run: pip install pre-commit
+      - name: Run pre-commit
+        run: pre-commit run --all-files

.pre-commit-config.yaml

@@ -0,0 +1,24 @@
+repos:
+  - repo: https://github.com/psf/black
+    rev: 24.4.0
+    hooks:
+      - id: black
+        args: ['--line-length=120', '--verbose']
+        exclude: '^models/'
+
+  - repo: https://github.com/pycqa/flake8
+    rev: '7.0.0'
+    hooks:
+      - id: flake8
+        exclude: '^models/'
+
+  - repo: https://github.com/pre-commit/mirrors-pylint
+    rev: v3.0.0a5
+    hooks:
+      - id: pylint
+        name: pylint
+        entry: pylint
+        language: system
+        args: ['.', '--rcfile=setup.cfg', '--fail-under=8']
+        exclude: '^models/'
+        types: [python]

01_📼_Upload_Video_File.py

@@ -6,44 +6,46 @@
 import moviepy.editor as mp
 import numpy as np
 import streamlit as st
-from streamlit_lottie import st_lottie
-from tqdm import tqdm
-
-from models.deep_colorization.colorizers import eccv16
-from utils import load_lottieurl, format_time, colorize_frame, change_model
 
-st.set_page_config(page_title="Image & Video Colorizer", page_icon="🎨", layout="wide")
+from tqdm import tqdm
 
-loaded_model = eccv16(pretrained=True).eval()
-current_model = "None"
+from utils import format_time, colorize_frame, change_model, load_model, setup_columns, set_page_config
 
-col1, col2 = st.columns([1, 3])
-with col1:
-    lottie = load_lottieurl("https://assets5.lottiefiles.com/packages/lf20_RHdEuzVfEL.json")
-    st_lottie(lottie)
+set_page_config()
+loaded_model = load_model()
+col2 = setup_columns()
+current_model = None
 
 with col2:
-    st.write("""
+    st.write(
+        """
     ## B&W Videos Colorizer
     ##### Upload a black and white video and get a colorized version of it.
     ###### ➠ This space is using CPU Basic so it might take a while to colorize a video.
-    ###### ➠ If you want more models and GPU available please support this space by donating.""")
+    ###### ➠ If you want more models and GPU available please support this space by donating."""
+    )
 
 
 def main():
+    """
+    Main function to run this page
+    """
     model = st.selectbox(
-        "Select Model (Both models have their pros and cons, I recommend trying both and keeping the best for your task)",
-        ["ECCV16", "SIGGRAPH17"], index=0)
+        "Select Model (Both models have their pros and cons, I recommend trying both and keeping the best for your "
+        "task)",
+        ["ECCV16", "SIGGRAPH17"],
+        index=0,
+    )
 
     loaded_model = change_model(current_model, model)
     st.write(f"Model is now {model}")
 
-    uploaded_file = st.file_uploader("Upload your video here...", type=['mp4', 'mov', 'avi', 'mkv'])
+    uploaded_file = st.file_uploader("Upload your video here...", type=["mp4", "mov", "avi", "mkv"])
 
     if st.button("Colorize"):
         if uploaded_file is not None:
             file_extension = os.path.splitext(uploaded_file.name)[1].lower()
-            if file_extension in ['.mp4', '.avi', '.mov', '.mkv']:
+            if file_extension in [".mp4", ".avi", ".mov", ".mkv"]:
                 # Save the video file to a temporary location
                 temp_file = tempfile.NamedTemporaryFile(delete=False)
                 temp_file.write(uploaded_file.read())
@@ -73,7 +75,7 @@ def main():
                         start_time = time.time()
                         time_text = st.text("Time Remaining: ")  # Initialize text value
 
-                        for _ in tqdm(range(total_frames), unit='frame', desc="Progress"):
+                        for _ in tqdm(range(total_frames), unit="frame", desc="Progress"):
                             ret, frame = video.read()
                             if not ret:
                                 break
@@ -123,7 +125,7 @@ def main():
                         st.download_button(
                             label="Download Colorized Video",
                             data=open(converted_filename, "rb").read(),
-                            file_name="colorized_video.mp4"
+                            file_name="colorized_video.mp4",
                         )
 
                         # Close and delete the temporary file after processing
@@ -135,4 +137,5 @@ def main():
     main()
     st.markdown(
         "###### Made with :heart: by [Clément Delteil](https://www.linkedin.com/in/clementdelteil/) [![this is an "
-        "image link](https://i.imgur.com/thJhzOO.png)](https://www.buymeacoffee.com/clementdelteil)")
+        "image link](https://i.imgur.com/thJhzOO.png)](https://www.buymeacoffee.com/clementdelteil)"
+    )

README.md

@@ -16,6 +16,18 @@ The following features are available:
 ## Interface
 <img src="interface.gif"/>
 
+## Running Locally
+If you want to process longer videos and you're limited by the Hugging Face space memory's limits, you can run this app locally.
+
+`ffmpeg.exe` is needed to run this app, you can install it using `brew install ffmpeg` and update the `IMAGEIO_FFMPEG_EXE` environment variable accordingly.
+
+```bash
+git clone https://github.com/Wazzabeee/image-video-colorization
+cd image-video-colorization
+pip install -r requirements.txt
+streamlit run 01_📼_Upload_Video_File.py
+```
+
 ## Todos
 Other models based on GANs will probably be implemented in the future if my application for a community grant to gain access to a GPU on Hugging Face is successful.
 

models/deep_colorization/__init__.py

@@ -0,0 +1,2 @@
+# inside models/deep_colorization/__init__.py
+from .colorizers import eccv16, siggraph17, load_img, preprocess_img, postprocess_tens

models/deep_colorization/colorizers/__init__.py

@@ -1,6 +1,4 @@
-
-from .base_color import *
-from .eccv16 import *
-from .siggraph17 import *
-from .util import *
-
+from .base_color import BaseColor
+from .eccv16 import ECCVGenerator, eccv16
+from .siggraph17 import SIGGRAPHGenerator, siggraph17
+from .util import load_img, resize_img, preprocess_img, postprocess_tens

models/deep_colorization/colorizers/base_color.py

@@ -1,24 +1,22 @@
-
-import torch
 from torch import nn
 
-class BaseColor(nn.Module):
-	def __init__(self):
-		super(BaseColor, self).__init__()
 
-		self.l_cent = 50.
-		self.l_norm = 100.
-		self.ab_norm = 110.
+class BaseColor(nn.Module):
+    def __init__(self):
+        super(BaseColor, self).__init__()
 
-	def normalize_l(self, in_l):
-		return (in_l-self.l_cent)/self.l_norm
+        self.l_cent = 50.0
+        self.l_norm = 100.0
+        self.ab_norm = 110.0
 
-	def unnormalize_l(self, in_l):
-		return in_l*self.l_norm + self.l_cent
+    def normalize_l(self, in_l):
+        return (in_l - self.l_cent) / self.l_norm
 
-	def normalize_ab(self, in_ab):
-		return in_ab/self.ab_norm
+    def unnormalize_l(self, in_l):
+        return in_l * self.l_norm + self.l_cent
 
-	def unnormalize_ab(self, in_ab):
-		return in_ab*self.ab_norm
+    def normalize_ab(self, in_ab):
+        return in_ab / self.ab_norm
 
+    def unnormalize_ab(self, in_ab):
+        return in_ab * self.ab_norm