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()