|
|
|
import torch |
|
from PIL import Image |
|
from RealESRGAN import RealESRGAN |
|
import gradio as gr |
|
import numpy as np |
|
import tempfile |
|
import time |
|
import zipfile |
|
import os |
|
|
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
|
def load_model(scale): |
|
model = RealESRGAN(device, scale=scale) |
|
weights_path = f'weights/RealESRGAN_x{scale}.pth' |
|
try: |
|
model.load_weights(weights_path, download=True) |
|
print(f"Weights for scale {scale} loaded successfully.") |
|
except Exception as e: |
|
print(f"Error loading weights for scale {scale}: {e}") |
|
model.load_weights(weights_path, download=False) |
|
return model |
|
|
|
|
|
model2 = load_model(2) |
|
model4 = load_model(4) |
|
model8 = load_model(8) |
|
|
|
def enhance_image(image, scale): |
|
try: |
|
print(f"Enhancing image with scale {scale}...") |
|
start_time = time.time() |
|
image_np = np.array(image.convert('RGB')) |
|
print(f"Image converted to numpy array: shape {image_np.shape}, dtype {image_np.dtype}") |
|
|
|
if scale == '2x': |
|
result = model2.predict(image_np) |
|
elif scale == '4x': |
|
result = model4.predict(image_np) |
|
else: |
|
result = model8.predict(image_np) |
|
|
|
enhanced_image = Image.fromarray(np.uint8(result)) |
|
print(f"Image enhanced in {time.time() - start_time:.2f} seconds") |
|
return enhanced_image |
|
except Exception as e: |
|
print(f"Error enhancing image: {e}") |
|
return image |
|
|
|
def muda_dpi(input_image, dpi): |
|
dpi_tuple = (dpi, dpi) |
|
image = Image.fromarray(input_image.astype('uint8'), 'RGB') |
|
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') |
|
image.save(temp_file, format='JPEG', dpi=dpi_tuple) |
|
temp_file.close() |
|
return Image.open(temp_file.name) |
|
|
|
def resize_image(input_image, width, height): |
|
image = Image.fromarray(input_image.astype('uint8'), 'RGB') |
|
resized_image = image.resize((width, height)) |
|
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') |
|
resized_image.save(temp_file, format='JPEG') |
|
temp_file.close() |
|
return Image.open(temp_file.name) |
|
|
|
def process_images(image_files, enhance, scale, adjust_dpi, dpi, resize, width, height): |
|
processed_images = [] |
|
temp_dir = tempfile.mkdtemp() |
|
|
|
for image_file in image_files: |
|
input_image = np.array(Image.open(image_file).convert('RGB')) |
|
original_image = Image.fromarray(input_image.astype('uint8'), 'RGB') |
|
|
|
if enhance: |
|
original_image = enhance_image(original_image, scale) |
|
|
|
if adjust_dpi: |
|
original_image = muda_dpi(np.array(original_image), dpi) |
|
|
|
if resize: |
|
original_image = resize_image(np.array(original_image), width, height) |
|
|
|
|
|
file_name = os.path.basename(image_file.name) |
|
output_path = os.path.join(temp_dir, file_name) |
|
original_image.save(output_path, format='JPEG') |
|
processed_images.append(output_path) |
|
|
|
|
|
zip_path = os.path.join(temp_dir, 'processed_images.zip') |
|
with zipfile.ZipFile(zip_path, 'w') as zipf: |
|
for file_path in processed_images: |
|
zipf.write(file_path, os.path.basename(file_path)) |
|
|
|
|
|
display_images = [Image.open(img_path) for img_path in processed_images] |
|
|
|
return display_images, zip_path |
|
|
|
iface = gr.Interface( |
|
fn=process_images, |
|
inputs=[ |
|
gr.Files(label="Upload Image Files"), |
|
gr.Checkbox(label="Enhance Images (ESRGAN)"), |
|
gr.Radio(['2x', '4x', '8x'], type="value", value='2x', label='Resolution model'), |
|
gr.Checkbox(label="Adjust DPI"), |
|
gr.Number(label="DPI", value=300), |
|
gr.Checkbox(label="Resize"), |
|
gr.Number(label="Width", value=512), |
|
gr.Number(label="Height", value=512) |
|
], |
|
outputs=[ |
|
gr.Gallery(label="Final Images"), |
|
gr.File(label="Download Final Images (ZIP)") |
|
], |
|
title="Multi-Image Enhancer", |
|
description="Upload multiple images (.jpg, .png), enhance using AI, adjust DPI, resize, and download the final results as a ZIP file." |
|
) |
|
|
|
iface.launch(debug=True) |
|
|
|
|
|
|
|
|