image-video-colorization

Runtime error

File size: 8,707 Bytes

import streamlit as st
import cv2
import numpy
import os
import random
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils.download_util import load_file_from_url
from PIL import Image

from realesrgan import RealESRGANer
from realesrgan.archs.srvgg_arch import SRVGGNetCompact


last_file = None
img_mode = "RGBA"


def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
    """Real-ESRGAN function to restore (and upscale) images.
    """
    if not img:
        return

    # Define model parameters
    if model_name == 'RealESRGAN_x4plus':  # x4 RRDBNet model
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        netscale = 4
        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
    elif model_name == 'RealESRNet_x4plus':  # x4 RRDBNet model
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        netscale = 4
        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth']
    elif model_name == 'RealESRGAN_x4plus_anime_6B':  # x4 RRDBNet model with 6 blocks
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
        netscale = 4
        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
    elif model_name == 'RealESRGAN_x2plus':  # x2 RRDBNet model
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
        netscale = 2
        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
    elif model_name == 'realesr-general-x4v3':  # x4 VGG-style model (S size)
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
        netscale = 4
        file_url = [
            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth',
            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'
        ]

    # Determine model paths
    model_path = os.path.join('weights', model_name + '.pth')
    if not os.path.isfile(model_path):
        ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
        for url in file_url:
            # model_path will be updated
            model_path = load_file_from_url(
                url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)

    # Use dni to control the denoise strength
    dni_weight = None
    if model_name == 'realesr-general-x4v3' and denoise_strength != 1:
        wdn_model_path = model_path.replace('realesr-general-x4v3', 'realesr-general-wdn-x4v3')
        model_path = [model_path, wdn_model_path]
        dni_weight = [denoise_strength, 1 - denoise_strength]

    # Restorer Class
    upsampler = RealESRGANer(
        scale=netscale,
        model_path=model_path,
        dni_weight=dni_weight,
        model=model,
        tile=0,
        tile_pad=10,
        pre_pad=10,
        half=False,
        gpu_id=None
    )

    # Use GFPGAN for face enhancement
    if face_enhance:
        from gfpgan import GFPGANer
        face_enhancer = GFPGANer(
            model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
            upscale=outscale,
            arch='clean',
            channel_multiplier=2,
            bg_upsampler=upsampler)

    # Convert the input PIL image to cv2 image, so that it can be processed by realesrgan
    #cv_img = numpy.array(img.get_value(), dtype = 'uint8')
    cv_img = numpy.array(img)
    #img = cv2.cvtColor(cv2.UMat(imgUMat), cv2.COLOR_RGB2GRAY)
    img = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)

    # Apply restoration
    try:
        if face_enhance:
            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
        else:
            output, _ = upsampler.enhance(img, outscale=outscale)
    except RuntimeError as error:
        print('Error', error)
        print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
    else:
        # Save restored image and return it to the output Image component
        if img_mode == 'RGBA':  # RGBA images should be saved in png format
            extension = 'png'
        else:
            extension = 'jpg'

        out_filename = f"output_{rnd_string(8)}.{extension}"
        cv2.imwrite(out_filename, output)
        global last_file
        last_file = out_filename
        return out_filename


def rnd_string(x):
    """Returns a string of 'x' random characters
    """
    characters = "abcdefghijklmnopqrstuvwxyz_0123456789"
    result = "".join((random.choice(characters)) for i in range(x))
    return result


def reset():
    """Resets the Image components of the Gradio interface and deletes
    the last processed image
    """
    global last_file
    if last_file:
        print(f"Deleting {last_file} ...")
        os.remove(last_file)
        last_file = None
    return gr.update(value=None), gr.update(value=None)


def has_transparency(img):
    """This function works by first checking to see if a "transparency" property is defined
    in the image's info -- if so, we return "True". Then, if the image is using indexed colors
    (such as in GIFs), it gets the index of the transparent color in the palette
    (img.info.get("transparency", -1)) and checks if it's used anywhere in the canvas
    (img.getcolors()). If the image is in RGBA mode, then presumably it has transparency in
    it, but it double-checks by getting the minimum and maximum values of every color channel
    (img.getextrema()), and checks if the alpha channel's smallest value falls below 255.
    https://stackoverflow.com/questions/43864101/python-pil-check-if-image-is-transparent
    """
    if img.info.get("transparency", None) is not None:
        return True
    if img.mode == "P":
        transparent = img.info.get("transparency", -1)
        for _, index in img.getcolors():
            if index == transparent:
                return True
    elif img.mode == "RGBA":
        extrema = img.getextrema()
        if extrema[3][0] < 255:
            return True
    return False


def image_properties(img):
    """Returns the dimensions (width and height) and color mode of the input image and
    also sets the global img_mode variable to be used by the realesrgan function
    """
    global img_mode
    if img:
        if has_transparency(img):
            img_mode = "RGBA"
        else:
            img_mode = "RGB"
        properties = f"Width: {img.size[0]}, Height: {img.size[1]}  |  Color Mode: {img_mode}"
        return properties
    
def image_properties(image):
    # Function to display image properties
    properties = f"Image Size: {image.size}\nImage Mode: {image.mode}"
    return properties

#----------

input_folder = '.'

@st.cache_resource
def load_image(image_file):
  img = Image.open(image_file)
  return img

def save_image(image_file):
  if image_file is not None:
    filename = image_file.name
    img = load_image(image_file)
    st.image(image=img, width=None)
    with open(os.path.join(input_folder, filename), "wb") as f:
      f.write(image_file.getbuffer())
      st.success("Succesfully uploaded file for processing".format(filename))

#------------

st.title("Image Denoizer")
# Saving uploaded image in input folder for processing

#with st.expander("Options/Parameters"):   

input_img = st.file_uploader(
"Upload Image", type=['png', 'jpeg', 'jpg', 'webp'])
#save_image(input_img)

model_name = "realesr-general-x4v3"

denoise_strength = st.slider("Denoise Strength", 0.0, 1.0, 0.5)

outscale = 1

face_enhance = False

if input_img:
  print(input_img)
  input_img = Image.open(input_img)
  # Display image properties
  cols = st.columns(2)

  cols[0].image(input_img, 'Source Image')

  #input_properties = get_image_properties(input_img)
  #cols[1].write(input_properties)

  # Output placeholder
  output_img = st.empty()

# Input and output placeholders
input_img = input_img
output_img = st.empty()

# Buttons
restore = st.button('Restore')
reset = st.button('Reset')

# Restore clicked
if restore:
    if input_img is not None:
        output = realesrgan(input_img, model_name, denoise_strength, 
                            face_enhance, outscale)
        output_img.image(output, 'Restored Image')
    else:
        st.warning('Upload a file', icon="⚠️")

# Reset clicked    
if reset:
    output_img.empty()