fix color
Browse files- .gitignore +4 -0
- app.py +17 -9
- model_image_colorizer.py +0 -33
.gitignore
CHANGED
@@ -1,3 +1,7 @@
|
|
1 |
.idea/
|
2 |
__pycache__/
|
|
|
|
|
|
|
|
|
3 |
|
|
|
1 |
.idea/
|
2 |
__pycache__/
|
3 |
+
models/
|
4 |
+
flagged/
|
5 |
+
result_images/
|
6 |
+
|
7 |
|
app.py
CHANGED
@@ -1,29 +1,37 @@
|
|
1 |
import os
|
|
|
2 |
from pathlib import Path
|
3 |
|
4 |
import gradio as gr
|
5 |
from deoldify import device
|
6 |
from deoldify.device_id import DeviceId
|
7 |
-
from deoldify.
|
8 |
from huggingface_hub import snapshot_download
|
9 |
|
10 |
-
from model_image_colorizer import ImageFilter, ModelImageColorizer
|
11 |
-
|
12 |
os.system("pip freeze")
|
|
|
13 |
|
14 |
device.set(device=DeviceId.CPU)
|
15 |
|
16 |
REPO_ID = "leonelhs/deoldify"
|
17 |
-
MODEL_NAME = "ColorizeArtistic_gen"
|
18 |
|
|
|
19 |
snapshot_folder = snapshot_download(repo_id=REPO_ID)
|
20 |
-
|
21 |
-
|
22 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
23 |
|
24 |
|
25 |
def predict(image):
|
26 |
-
return colorizer.
|
27 |
|
28 |
|
29 |
title = "DeOldify"
|
@@ -47,7 +55,7 @@ This demo is running on a CPU, if you like this project please make us a donatio
|
|
47 |
|
48 |
demo = gr.Interface(
|
49 |
predict, [
|
50 |
-
gr.Image(type="
|
51 |
], [
|
52 |
gr.Image(type="pil", label="Image color")
|
53 |
],
|
|
|
1 |
import os
|
2 |
+
import warnings
|
3 |
from pathlib import Path
|
4 |
|
5 |
import gradio as gr
|
6 |
from deoldify import device
|
7 |
from deoldify.device_id import DeviceId
|
8 |
+
from deoldify.visualize import get_image_colorizer
|
9 |
from huggingface_hub import snapshot_download
|
10 |
|
|
|
|
|
11 |
os.system("pip freeze")
|
12 |
+
warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")
|
13 |
|
14 |
device.set(device=DeviceId.CPU)
|
15 |
|
16 |
REPO_ID = "leonelhs/deoldify"
|
17 |
+
MODEL_NAME = "ColorizeArtistic_gen.pth"
|
18 |
|
19 |
+
model_path = Path("models")
|
20 |
snapshot_folder = snapshot_download(repo_id=REPO_ID)
|
21 |
+
model_path_hf = Path(snapshot_folder).joinpath(MODEL_NAME)
|
22 |
+
|
23 |
+
Path.mkdir(model_path, exist_ok=True)
|
24 |
+
try:
|
25 |
+
Path.symlink_to(model_path.joinpath(MODEL_NAME), model_path_hf)
|
26 |
+
except FileExistsError:
|
27 |
+
pass
|
28 |
+
|
29 |
+
device.set(device=DeviceId.GPU0)
|
30 |
+
colorizer = get_image_colorizer(artistic=True)
|
31 |
|
32 |
|
33 |
def predict(image):
|
34 |
+
return colorizer.get_transformed_image(image, render_factor=35, watermarked=False)
|
35 |
|
36 |
|
37 |
title = "DeOldify"
|
|
|
55 |
|
56 |
demo = gr.Interface(
|
57 |
predict, [
|
58 |
+
gr.Image(type="filepath", label="Image gray scale"),
|
59 |
], [
|
60 |
gr.Image(type="pil", label="Image color")
|
61 |
],
|
model_image_colorizer.py
DELETED
@@ -1,33 +0,0 @@
|
|
1 |
-
import torch
|
2 |
-
from PIL import Image as PilImage
|
3 |
-
|
4 |
-
from deoldify.filters import IFilter, BaseFilter
|
5 |
-
from deoldify.visualize import ModelImageVisualizer
|
6 |
-
from fastai.basic_train import Learner
|
7 |
-
from fastai.vision import normalize_funcs
|
8 |
-
|
9 |
-
stats = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
10 |
-
|
11 |
-
|
12 |
-
class ImageFilter(BaseFilter):
|
13 |
-
def __init__(self, learn: Learner):
|
14 |
-
super().__init__(learn)
|
15 |
-
self.render_base = 16
|
16 |
-
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
17 |
-
self.norm, self.denorm = normalize_funcs(*stats)
|
18 |
-
|
19 |
-
def filter(self, filtered_image: PilImage, render_factor=35) -> PilImage:
|
20 |
-
orig_image = filtered_image.copy()
|
21 |
-
render_sz = render_factor * self.render_base
|
22 |
-
model_image = self._model_process(orig=filtered_image, sz=render_sz)
|
23 |
-
raw_color = self._unsquare(model_image, orig_image)
|
24 |
-
return raw_color
|
25 |
-
|
26 |
-
|
27 |
-
class ModelImageColorizer(ModelImageVisualizer):
|
28 |
-
def __init__(self, filter: IFilter):
|
29 |
-
self.filter = filter
|
30 |
-
|
31 |
-
def get_colored_image(self, image, render_factor: int = None) -> PilImage:
|
32 |
-
self._clean_mem()
|
33 |
-
return self.filter.filter(image, render_factor=render_factor)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|