File size: 2,136 Bytes
5e30dca
 
 
e9ea41c
5e30dca
 
 
e9ea41c
 
 
4c97d0f
 
 
3c40616
4c97d0f
 
 
 
3c40616
4c97d0f
 
 
3c40616
4c97d0f
 
 
 
3c40616
4c97d0f
e9ea41c
 
 
4c97d0f
e9ea41c
 
 
 
 
 
 
 
4c97d0f
e9ea41c
 
4c97d0f
5e30dca
e9ea41c
5e30dca
 
 
4c97d0f
5e30dca
 
4c97d0f
5e30dca
4c97d0f
5e30dca
4c97d0f
 
 
 
 
 
e9ea41c
4c97d0f
 
e9ea41c
4c97d0f
5e30dca
 
4c97d0f
 
3c40616
4c97d0f
5e30dca
 
 
e9ea41c
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
78
79
import os
import cv2
import torch
import gradio as gr
import numpy as np
from PIL import Image
from gfpgan import GFPGANer
from basicsr.archs.srvgg_arch import SRVGGNetCompact
from realesrgan.utils import RealESRGANer

# Download GFPGAN model if not already present
GFPGAN_MODEL = "GFPGANv1.4.pth"
GFPGAN_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"

if not os.path.exists(GFPGAN_MODEL):
    import requests
    r = requests.get(GFPGAN_URL, allow_redirects=True)
    open(GFPGAN_MODEL, 'wb').write(r.content)

# Download RealESRGAN model
ESRGAN_MODEL = "realesr-general-x4v3.pth"
ESRGAN_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"

if not os.path.exists(ESRGAN_MODEL):
    import requests
    r = requests.get(ESRGAN_URL, allow_redirects=True)
    open(ESRGAN_MODEL, 'wb').write(r.content)

# Initialize Real-ESRGAN as background upsampler
esr_model = SRVGGNetCompact(
    num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32,
    upscale=4, act_type='prelu'
)

bg_upsampler = RealESRGANer(
    scale=4,
    model_path=ESRGAN_MODEL,
    model=esr_model,
    tile=0,
    tile_pad=10,
    pre_pad=0,
    half=torch.cuda.is_available()
)

# Initialize GFPGAN with Real-ESRGAN as background upsampler
restorer = GFPGANer(
    model_path=GFPGAN_MODEL,
    upscale=2,
    arch='clean',
    channel_multiplier=2,
    bg_upsampler=bg_upsampler
)

# Enhancement function
def enhance(image):
    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

    _, _, restored_img = restorer.enhance(
        img,
        has_aligned=False,
        only_center_face=False,
        paste_back=True
    )

    restored_pil = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB))
    return restored_pil

# Gradio UI
iface = gr.Interface(
    fn=enhance,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="IMGEN - AI Photo Enhancer (Face + Outfit)",
    description="Upload your photo (ID, CV, profile) and enhance both the face and outfit with AI using GFPGAN + RealESRGAN."
)

if __name__ == "__main__":
    iface.launch()