import streamlit as st
import os.path

os.system("mkdir _input")
os.system("mkdir _output")
os.system("mkdir _outputf")
os.system("ls")
if not os.path.isfile("./_input/imagem-0001.png"):
    os.system("ffmpeg -i vivi.mp4   -compression_level 10 -pred mixed -pix_fmt rgb24 -sws_flags +accurate_rnd+full_chroma_int -s 1080x1920 -r 0.12 ./_input/imagem-%4d.png")

os.system("ls ./_input") 
if 'myVar' not in globals():
    myVar=""
    os.system("pip install git+https://github.com/TencentARC/GFPGAN.git")

import cv2
import glob
import numpy as np
from basicsr.utils import imwrite
from gfpgan import GFPGANer
#os.system("pip freeze")
#os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v0.2.0/GFPGANCleanv1-NoCE-C2.pth -P .")
import random

from PIL import Image
import torch
# torch.hub.download_url_to_file('https://upload.wikimedia.org/wikipedia/commons/thumb/a/ab/Abraham_Lincoln_O-77_matte_collodion_print.jpg/1024px-Abraham_Lincoln_O-77_matte_collodion_print.jpg', 'lincoln.jpg')
# torch.hub.download_url_to_file('https://upload.wikimedia.org/wikipedia/commons/5/50/Albert_Einstein_%28Nobel%29.png', 'einstein.png')
# torch.hub.download_url_to_file('https://upload.wikimedia.org/wikipedia/commons/thumb/9/9d/Thomas_Edison2.jpg/1024px-Thomas_Edison2.jpg', 'edison.jpg')
# torch.hub.download_url_to_file('https://upload.wikimedia.org/wikipedia/commons/thumb/a/a9/Henry_Ford_1888.jpg/1024px-Henry_Ford_1888.jpg', 'Henry.jpg')
# torch.hub.download_url_to_file('https://upload.wikimedia.org/wikipedia/commons/thumb/0/06/Frida_Kahlo%2C_by_Guillermo_Kahlo.jpg/800px-Frida_Kahlo%2C_by_Guillermo_Kahlo.jpg', 'Frida.jpg')

# set up GFPGAN restorer
bg_upsampler = None
print(f"Is CUDA available: {torch.cuda.is_available()}")
if 'restorer' not in globals():
    restorer = GFPGANer(
        model_path='GFPGANv1.3.pth',
        upscale=2,
        arch='clean',
        channel_multiplier=2,
        bg_upsampler=bg_upsampler)


img_list = sorted(glob.glob(os.path.join("./_input", '*')))

for img_path in img_list:
        # read image
        img_name = os.path.basename(img_path)
        print(f'Processing {img_name} ...')
        basename, ext = os.path.splitext(img_name)
        input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)

        # restore faces and background if necessary
        cropped_faces, restored_faces, restored_img = restorer.enhance(
            input_img,
            has_aligned='store_true',
            only_center_face='store_true',
            paste_back=True,
            weight=0.5)

        # save faces
        for idx, (cropped_face, restored_face) in enumerate(zip(cropped_faces, restored_faces)):
            # save cropped face
            save_crop_path = os.path.join("_output", 'cropped_faces', f'{basename}_{idx:02d}.png')
            imwrite(cropped_face, save_crop_path)
            # save restored face
            if None is not None:
                save_face_name = f'{basename}_{idx:04d}_{args.suffix}.png'
            else:
                save_face_name = f'{basename}_{idx:04d}.png'
            save_restore_path = os.path.join("_output", 'restored_faces', save_face_name)
            imwrite(restored_face, save_restore_path)
            # save comparison image
            cmp_img = np.concatenate((cropped_face, restored_face), axis=1)
            imwrite(cmp_img, os.path.join("_output", 'cmp', f'{basename}_{idx:04d}.png'))

        # save restored img
        if restored_img is not None:
            print('encontrou**************')          
            if args.ext == 'auto':
                extension = ext[1:]
            else:
                extension = args.ext

            if None is not None:
                save_restore_path = os.path.join("_output", 'restored_imgs', f'{basename}_{args.suffix}.{extension}')
            else:
                save_restore_path = os.path.join("_output", 'restored_imgs', f'{basename}.{extension}')
            imwrite(restored_img, save_restore_path)
os.system("ls ./_output")
os.system("echo ----")
os.system("ls ./_output/cmp")
os.system("echo ----")
os.system("ls ./_output/restored_imgs")
os.system("echo ----")




def inference():
    random.randint(0, 9)
    input_img = cv2.imread("./_output/cmp/imagem-000"+str(random.randint(1, 4))+"_0000.png" , cv2.IMREAD_COLOR)
    input_img= cv2.cvtColor(input_img,cv2.COLOR_BGR2RGB) 
    st.image(input_img)
   
    #return Image.fromarray(restored_faces[0][:,:,::-1])
    

title = "Melhoria de imagens"

os.system("ls")
description = "Sistema para automação。"

article = "<p style='text-align: center'><a href='https://huggingface.co/spaces/akhaliq/GFPGAN/' target='_blank'>clone from akhaliq@huggingface with little change</a> | <a href='https://github.com/TencentARC/GFPGAN' target='_blank'>GFPGAN Github Repo</a></p><center><img src='https://visitor-badge.glitch.me/badge?page_id=akhaliq_GFPGAN' alt='visitor badge'></center>"
st.button('Comparacao',on_click=inference)