File size: 3,141 Bytes
6ac308a |
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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
"""
Enhance Utils
-------------
Utility functions untuk meningkatkan kualitas gambar dengan menggunakan model AI:
- Real-ESRGAN untuk upscaling gambar umum
- GFPGAN untuk enhancement wajah (face restoration)
Created by _drat | 2025
"""
import os
import torch
import cv2
import numpy as np
import subprocess
from PIL import Image
from gfpgan.utils import GFPGANer
from basicsr.archs.srvgg_arch import SRVGGNetCompact
from realesrgan.utils import RealESRGANer
# Fungsi untuk menjalankan perintah shell (misal: pip, wget)
def runcmd(cmd, verbose = False, *args, **kwargs):
process = subprocess.Popen(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.PIPE,
text = True,
shell = True
)
std_out, std_err = process.communicate()
if verbose:
print(std_out.strip(), std_err)
pass
# Debug: Menampilkan daftar library yang terinstall di environment
runcmd("pip freeze")
# Download model GFPGAN jika belum ada di direktori saat ini
if not os.path.exists('GFPGANv1.4.pth'):
runcmd("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
# Download model RealESRGAN jika belum ada di direktori saat ini
if not os.path.exists('realesr-general-x4v3.pth'):
runcmd("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
# Inisialisasi model Real-ESRGAN (untuk upscaling gambar umum)
model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
model_path = 'realesr-general-x4v3.pth'
half = True if torch.cuda.is_available() else False # Pakai half precision jika tersedia GPU
upsampler = RealESRGANer(
scale=4,
model_path=model_path,
model=model,
tile=0,
tile_pad=10,
pre_pad=0,
half=half
)
# Inisialisasi enhancer wajah (GFPGAN)
face_enhancer = GFPGANer(
model_path='GFPGANv1.4.pth',
upscale=1, # Default tidak melakukan upscaling (hanya enhance wajah)
arch='clean',
channel_multiplier=2
)
# Fungsi utama enhancement gambar
def enhance_image(
pil_image: Image, # Input gambar dalam format PIL
enhance_face: bool = False, # Apakah perlu enhance wajah?
):
# Konversi PIL image ke array BGR (OpenCV)
img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
h, w = img.shape[0:2]
# Optional: jika gambar kecil, perbesar dulu agar hasil enhancement lebih maksimal
if h < 300:
img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
# Enhance menggunakan GFPGAN (wajah) atau Real-ESRGAN (umum)
if enhance_face:
# Hanya enhance bagian wajah utama di tengah gambar
_, _, output = face_enhancer.enhance(
img, has_aligned=False, only_center_face=True, paste_back=True
)
else:
# Upscale seluruh gambar dengan Real-ESRGAN (outscale=2x)
output, _ = upsampler.enhance(img, outscale=2)
# Konversi hasil output ke format PIL RGB untuk siap digunakan di aplikasi
pil_output = Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB))
return pil_output
|