110 lines
3.3 KiB
Python
110 lines
3.3 KiB
Python
import modules.scripts as scripts
|
|
import gradio as gr
|
|
import numpy as np
|
|
import cv2 as cv
|
|
|
|
from modules.processing import process_images, get_fixed_seed
|
|
from copy import copy
|
|
|
|
# https://docs.opencv.org/4.8.0/d2/df0/tutorial_py_hdr.html
|
|
def merge_HDR(imgs:list, path:str, depth:str, fmt:str, gamma:float):
|
|
import datetime
|
|
import math
|
|
import os
|
|
|
|
output_folder = os.path.join(path, 'hdr')
|
|
if not os.path.exists(output_folder):
|
|
os.makedirs(output_folder)
|
|
|
|
imgs_np = [np.array(img, dtype=np.uint8) for img in imgs]
|
|
|
|
merge = cv.createMergeMertens()
|
|
hdr = merge.process(imgs_np)
|
|
hdr += math.ceil(0 - np.min(hdr) * 1000) / 1000
|
|
|
|
#print(f'{np.min(hdr)}, {np.max(hdr)}')
|
|
|
|
target = (65535 if depth == '16bpc' else 255)
|
|
precision = ('uint16' if depth == '16bpc' else 'uint8')
|
|
|
|
hdr = np.power(hdr, (1 / gamma))
|
|
|
|
ldr = np.clip(hdr * target, 0, target).astype(precision)
|
|
rgb = cv.cvtColor(ldr, cv.COLOR_BGR2RGB)
|
|
|
|
cv.imwrite(os.path.join(output_folder, f'{datetime.datetime.now().strftime("%H-%M-%S")}{fmt}'), rgb)
|
|
|
|
class VectorHDR(scripts.Script):
|
|
|
|
def title(self):
|
|
return "High Dynamic Range"
|
|
|
|
def show(self, is_img2img):
|
|
return True
|
|
|
|
def ui(self, is_img2img):
|
|
with gr.Row():
|
|
count = gr.Slider(label="Brackets", minimum=3, maximum=9, step=2, value=7)
|
|
gap = gr.Slider(label="Gaps", minimum=0.50, maximum=2.50, step=0.25, value=1.50)
|
|
|
|
with gr.Accordion("Merge Options", elem_id='vec-hdr-' + ('img' if is_img2img else 'txt'), open=False):
|
|
auto = gr.Checkbox(label="Automatically Merge", value=True)
|
|
|
|
with gr.Row():
|
|
depth = gr.Radio(['16bpc', '8bpc'], label="Bit Depth", value="16bpc")
|
|
fmt = gr.Radio(['.tiff', '.png'], label="Image Format", value=".tiff")
|
|
|
|
gamma = gr.Slider(label="Gamma", info='Lower: Darker | Higher: Brighter',minimum=0.2, maximum=2.2, step=0.2, value=1.2)
|
|
|
|
return [count, gap, auto, depth, fmt, gamma]
|
|
|
|
def run(self, p, count:int, gap:float, auto:bool, depth:str, fmt:str, gamma:float):
|
|
center = count // 2
|
|
|
|
p.seed = get_fixed_seed(p.seed)
|
|
p.scripts.script('vectorscope cc').xyzCache.update({
|
|
'Enable':'True',
|
|
'Alt':'True',
|
|
'Skip': 0,
|
|
'Brightness': 0,
|
|
'DoHR': 'False',
|
|
'Method': 'Ones',
|
|
'Scaling': '1 - Cos'
|
|
})
|
|
|
|
baseline = process_images(p)
|
|
pc = copy(p)
|
|
|
|
imgs = [None] * count
|
|
imgs[center] = baseline.images[0]
|
|
|
|
brackets = brightness_brackets(count, gap)
|
|
|
|
for it in range(count):
|
|
if it == center:
|
|
continue
|
|
|
|
pc.scripts.script('vectorscope cc').xyzCache.update({
|
|
'Enable':'True',
|
|
'Alt':'True',
|
|
'Skip': 0,
|
|
'Brightness': brackets[it],
|
|
'DoHR': 'False',
|
|
'Method': 'Ones',
|
|
'Scaling': '1 - Cos'
|
|
})
|
|
|
|
proc = process_images(pc)
|
|
imgs[it] = proc.images[0]
|
|
|
|
if not auto:
|
|
baseline.images = imgs
|
|
return baseline
|
|
else:
|
|
merge_HDR(imgs, p.outpath_samples, depth, fmt, gamma)
|
|
return baseline
|
|
|
|
def brightness_brackets(count, gap):
|
|
half = count // 2
|
|
return [gap * (i - half) for i in range(count)]
|