File size: 1,971 Bytes
08e633b 0a07ad3 08e633b 0a07ad3 08e633b 0a07ad3 ddd3994 ede4c6e 0d9b367 ede4c6e 08e633b 0a07ad3 20f44f5 0a07ad3 20f44f5 0a07ad3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import os
import warnings
from pathlib import Path
import gradio as gr
from gradio_imageslider import ImageSlider
import numpy as np
import torch
import fastai
from deoldify import device
from deoldify.device_id import DeviceId
from deoldify.visualize import *
from huggingface_hub import snapshot_download
os.system("pip freeze")
from collections.abc import Sized # Import Sized from collections.abc
warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")
# private repo
from dotenv import load_dotenv
load_dotenv()
token = os.getenv("HF_TOKEN")
snapshot_folder = snapshot_download(repo_id="afondiel/image-colorizer-deoldify", token=token)
# Set the device to use for computation
# choices: CPU, GPU0...GPU7
if torch.cuda.is_available():
device.set(device=DeviceId.GPU0)
else:
device.set(device=DeviceId.CPU)
# Load the pre-trained model
_colorizer = get_image_colorizer(root_folder=Path(snapshot_folder), artistic=True)
def colorizer_fn(input_img, render_factor):
"""
Colorize grayscale images/photos
- @param input_img old (grayscale) image
- @param render_factor render_factor
"""
if input_img is not None and input_img !='':
output_img = _colorizer.get_transformed_image(
path=input_img,
render_factor=int(render_factor),
watermarked=watermarked,
post_process=True,
)
else:
print('Provide an image and try again.')
return (input_img, output_img) # Return a tuple of old and color Image to be plotted with ImageSlider()
title = "AI Image Colorizer"
description = "Colorize old images with AI"
examples = [["./demo.jpg"],]
demo = gr.Interface(
fn=colorizer_fn,
inputs=[gr.Image(type="filepath" , label="Old image"), gr.Slider(0, 40, label="Render Factor", value=10)],
outputs=ImageSlider(type="pil", label="Old vs Colored image"),
examples=examples,
title=title,
description=description,
)
# Launch the demo
demo.launch()
|