from PIL import Image, ImageDraw, ImageFont, ImageEnhance
import torch
import torchvision
import gradio
import numpy as np

Instructuction = "Select a Unique Portrait Image of yourself"
title="I am something of a Painter, Anime-Edition (AnimeGAN-V2)"
description = "Drag/Drop or Upload a cute portrait Image of yourself or anyone you find interesting 😉, then observe how this Generative model\
               is able to Generate a cute Anime-Cartoon version of your Image."
article = """
            - Select an image from the examples provided as demo image,
            - Click submit to Generate Image,
            - Tips: Quality Images with Great brightness/without pre-existing filters work better (Image Noise).
            - Privacy: No user data is collected or saved,
            - Credits to akhaliq/AnimeGANv2 for original AnimeGanV2"""

model = torch.hub.load(
    "AK391/animegan2-pytorch:main",
    "generator",
    pretrained=True,
    #device="cuda",
    progress=False
)
face2paint = torch.hub.load(
    'AK391/animegan2-pytorch:main', 'face2paint', 
    size=512, #device="cuda",
    side_by_side=False
)

def enhance_logo_brightness(logo_image, factor):
    enhancer = ImageEnhance.Brightness(logo_image)
    return enhancer.enhance(factor)

def add_logo(image):
    logo_path = "./logo.jpg" # 로고 파일 경로
    logo_size = (100, 30) # 로고 크기
    logo = Image.open(logo_path).convert("RGBA")
    logo = logo.resize(logo_size, resample=Image.BICUBIC)
    logo = enhance_logo_brightness(logo, 1.5) # 로고 이미지의 밝기를 조절합니다.
    return logo

def inference(img):
    img_pil = Image.fromarray(img.astype('uint8'), 'RGB') # 이미지를 PIL 이미지 객체로 변환
    output_image_pil = face2paint(model, img_pil.resize((512, 512))) # 이미지 변환
    output_image_pil = output_image_pil.convert('RGB') # RGBA -> RGB 변환
    logo = add_logo(output_image_pil) # 로고 이미지 가져오기
    output_image_pil.paste(logo, (512 - logo.size[0], 512 - logo.size[1]), logo) # 로고 이미지를 생성된 이미지에 붙이기
    output_image = np.array(output_image_pil) # PIL 이미지 객체를 numpy 배열로 변환
    return output_image

gradio.Interface(inference,
                 inputs=gradio.Image(type='numpy'),
                 outputs=gradio.Image(type='numpy'),
                 Instructuction=Instructuction, title=title, description=description, article=article).launch()