File size: 2,222 Bytes
08e633b 0a07ad3 78e67c7 0e8d222 0a07ad3 08e633b 0a07ad3 78e67c7 0a07ad3 78e67c7 0d9b367 cf8ef53 af1a89a cf8ef53 0e8d222 cf8ef53 0e8d222 0a07ad3 78e67c7 0a07ad3 20f44f5 0a07ad3 0e8d222 0a07ad3 78e67c7 0a07ad3 78e67c7 0a07ad3 78e67c7 0a07ad3 af1a89a 78e67c7 |
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 72 73 74 75 76 77 |
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 HfApi, snapshot_download
from huggingface_hub import HfApi, HfFolder
os.system("pip freeze")
from collections.abc import Sized # Import Sized from collections.abc
# Suppress warnings
warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")
# Private repo settings
token = os.getenv("HF_TOKEN")
repo_id = "afondiel/image-colorizer-deoldify"
repo_type = "space"
device.set(device=DeviceId.CPU)
# Initialize Hugging Face API
api = HfApi(token=token)
HfFolder.save_token(token)
# Download the snapshot from the space repository
snapshot_folder = api.snapshot_download(repo_id=repo_id, repo_type=repo_type, revision="main", token=token)
device.set(device=DeviceId.GPU0)
# Load the pre-trained model
_colorizer = get_image_colorizer(root_folder=Path(snapshot_folder), artistic=True)
print(f"root_folder: {root_folder}")
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
if __name__ == "__main__":
demo.launch()
|