Enchance images in a directory using EnlightenGAN

generate_englighten.py

import os

import cv2
from enlighten_inference import EnlightenOnnxModel
from tqdm import tqdm

# Set input and output path
input_path = 'image_rgb'
output_path = 'image_enlighten'

# Create output folder
os.makedirs(output_path, exist_ok=True)

# Get file list
filelist = os.listdir(input_path)
filelist.sort()

# Load model
model = EnlightenOnnxModel()

# Process images
for filename in tqdm(filelist):
    img = cv2.imread(os.path.join(input_path, filename))
    processed = model.predict(img)
    cv2.imwrite(os.path.join(output_path, filename), processed)