leonelhs commited on
Commit
56ecbb4
1 Parent(s): e0cd300
Files changed (4) hide show
  1. .gitignore +3 -0
  2. app.py +58 -0
  3. model_image_colorizer.py +33 -0
  4. requirements.txt +9 -0
.gitignore ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ .idea/
2
+ __pycache__/
3
+
app.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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.generators import gen_inference_deep
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
+ learn = gen_inference_deep(root_folder=Path(snapshot_folder), weights_name=MODEL_NAME)
21
+ image_filter = ImageFilter(learn=learn)
22
+ colorizer = ModelImageColorizer(image_filter)
23
+
24
+
25
+ def inference(image):
26
+ return colorizer.get_colored_image(image, render_factor=35)
27
+
28
+
29
+ title = "DeOldify"
30
+ description = r"""
31
+ ## Colorize image
32
+
33
+ This is an implementation of <a href='https://github.com/jantic/DeOldify' target='_blank'>DeOldify</a>.
34
+ It has no any particular purpose than start research on AI models.
35
+
36
+ """
37
+
38
+ article = r"""
39
+ Questions, doubts, comments, please email 📧 `leonelhs@gmail.com`
40
+
41
+ This demo is running on a CPU, if you like this project please make us a donation to run on a GPU or just give us a <a href='https://github.com/jantic/DeOldify' target='_blank'>Github ⭐</a>
42
+
43
+ <a href="https://www.buymeacoffee.com/leonelhs"><img src="https://img.buymeacoffee.com/button-api/?text=Buy me a coffee&emoji=&slug=leonelhs&button_colour=FFDD00&font_colour=000000&font_family=Cookie&outline_colour=000000&coffee_colour=ffffff" /></a>
44
+
45
+ <center><img src='https://visitor-badge.glitch.me/badge?page_id=deoldify.visitor-badge' alt='visitor badge'></center>
46
+ """
47
+
48
+ demo = gr.Interface(
49
+ inference, [
50
+ gr.Image(type="pil", label="Image gray scale"),
51
+ ], [
52
+ gr.Image(type="pil", label="Image color")
53
+ ],
54
+ title=title,
55
+ description=description,
56
+ article=article)
57
+
58
+ demo.queue().launch()
model_image_colorizer.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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)
requirements.txt ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ torch>=2.0.1
2
+ torchvision~=0.15.2
3
+ numpy~=1.24.3
4
+ pillow~=9.5.0
5
+ opencv-python
6
+ deoldify~=0.0.1
7
+
8
+
9
+