fix color

.gitignore

@@ -1,3 +1,7 @@
 .idea/
 __pycache__/
+models/
+flagged/
+result_images/
+
 

app.py

@@ -1,29 +1,37 @@
 import os
+import warnings
 from pathlib import Path
 
 import gradio as gr
 from deoldify import device
 from deoldify.device_id import DeviceId
-from deoldify.generators import gen_inference_deep
+from deoldify.visualize import get_image_colorizer
 from huggingface_hub import snapshot_download
 
-from model_image_colorizer import ImageFilter, ModelImageColorizer
-
 os.system("pip freeze")
+warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")
 
 device.set(device=DeviceId.CPU)
 
 REPO_ID = "leonelhs/deoldify"
-MODEL_NAME = "ColorizeArtistic_gen"
+MODEL_NAME = "ColorizeArtistic_gen.pth"
 
+model_path = Path("models")
 snapshot_folder = snapshot_download(repo_id=REPO_ID)
-learn = gen_inference_deep(root_folder=Path(snapshot_folder), weights_name=MODEL_NAME)
-image_filter = ImageFilter(learn=learn)
-colorizer = ModelImageColorizer(image_filter)
+model_path_hf = Path(snapshot_folder).joinpath(MODEL_NAME)
+
+Path.mkdir(model_path, exist_ok=True)
+try:
+    Path.symlink_to(model_path.joinpath(MODEL_NAME), model_path_hf)
+except FileExistsError:
+    pass
+
+device.set(device=DeviceId.GPU0)
+colorizer = get_image_colorizer(artistic=True)
 
 
 def predict(image):
-    return colorizer.get_colored_image(image, render_factor=35)
+    return colorizer.get_transformed_image(image, render_factor=35, watermarked=False)
 
 
 title = "DeOldify"
@@ -47,7 +55,7 @@ This demo is running on a CPU, if you like this project please make us a donatio
 
 demo = gr.Interface(
     predict, [
-        gr.Image(type="pil", label="Image gray scale"),
+        gr.Image(type="filepath", label="Image gray scale"),
     ], [
         gr.Image(type="pil", label="Image color")
     ],

model_image_colorizer.py

@@ -1,33 +0,0 @@
-import torch
-from PIL import Image as PilImage
-
-from deoldify.filters import IFilter, BaseFilter
-from deoldify.visualize import ModelImageVisualizer
-from fastai.basic_train import Learner
-from fastai.vision import normalize_funcs
-
-stats = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-
-
-class ImageFilter(BaseFilter):
-    def __init__(self, learn: Learner):
-        super().__init__(learn)
-        self.render_base = 16
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        self.norm, self.denorm = normalize_funcs(*stats)
-
-    def filter(self, filtered_image: PilImage, render_factor=35) -> PilImage:
-        orig_image = filtered_image.copy()
-        render_sz = render_factor * self.render_base
-        model_image = self._model_process(orig=filtered_image, sz=render_sz)
-        raw_color = self._unsquare(model_image, orig_image)
-        return raw_color
-
-
-class ModelImageColorizer(ModelImageVisualizer):
-    def __init__(self, filter: IFilter):
-        self.filter = filter
-
-    def get_colored_image(self, image, render_factor: int = None) -> PilImage:
-        self._clean_mem()
-        return self.filter.filter(image, render_factor=render_factor)