Spaces:
Running
Running
import os | |
import time | |
import logging | |
import numpy as np | |
import keras.optimizers | |
import cv2 | |
import gradio as gr | |
import Network | |
import utls | |
class ImageProcessor: | |
def __init__(self, model_name): | |
self.model_name = model_name | |
self.mbllen = Network.build_mbllen((None, None, 3)) | |
self.mbllen.load_weights(f'{model_name}.h5') | |
self.opt = keras.optimizers.Adam(lr=2 * 1e-04, beta_1=0.9, beta_2=0.999, epsilon=1e-08) | |
self.mbllen.compile(loss='mse', optimizer=self.opt) | |
def process_image(self, image, highpercent, lowpercent, gamma, maxrange): | |
img_A = np.array(image) / 255.0 | |
img_A = img_A[np.newaxis, :] | |
maxrange /= 10. | |
hsvgamma = gamma / 10. | |
out_pred = self.mbllen.predict(img_A) | |
fake_B = out_pred[0, :, :, :3] | |
fake_B = self.adjust_image(fake_B, maxrange, highpercent, lowpercent, hsvgamma) | |
return fake_B | |
def adjust_image(self, fake_B, maxrange, highpercent, lowpercent, hsvgamma): | |
gray_fake_B = fake_B[:, :, 0] * 0.299 + fake_B[:, :, 1] * 0.587 + fake_B[:, :, 1] * 0.114 | |
percent_max = sum(sum(gray_fake_B >= maxrange)) / sum(sum(gray_fake_B <= 1.0)) | |
max_value = np.percentile(gray_fake_B[:], highpercent) | |
if percent_max < (100-highpercent)/100.: | |
scale = maxrange / max_value | |
fake_B = fake_B * scale | |
fake_B = np.minimum(fake_B, 1.0) | |
gray_fake_B = fake_B[:,:,0]*0.299 + fake_B[:,:,1]*0.587 + fake_B[:,:,1]*0.114 | |
sub_value = np.percentile(gray_fake_B[:], lowpercent) | |
fake_B = (fake_B - sub_value) * (1./(1-sub_value)) | |
imgHSV = cv2.cvtColor(fake_B, cv2.COLOR_RGB2HSV) | |
H, S, V = cv2.split(imgHSV) | |
S = np.power(S, hsvgamma) | |
imgHSV = cv2.merge([H, S, V]) | |
fake_B = cv2.cvtColor(imgHSV, cv2.COLOR_HSV2RGB) | |
fake_B = np.minimum(fake_B, 1.0) | |
return fake_B | |
if __name__ == "__main__": | |
logging.basicConfig(filename='image_processor.log', level=logging.INFO) | |
image_processor = ImageProcessor(model_name='Syn_img_lowlight_withnoise') | |
iface = gr.Interface( | |
image_processor.process_image, | |
[ | |
gr.inputs.Image(shape=(None, None)), # Corrected line | |
gr.inputs.Slider(85, 100, default=95, label="High Percent"), | |
gr.inputs.Slider(0, 100, default=5, label="Low Percent"), | |
gr.inputs.Slider(6, 10, default=8, label="Gamma"), | |
gr.inputs.Slider(0, 20, default=8, label="Max Range"), | |
], | |
gr.outputs.Image(type="numpy"), | |
) | |
iface.launch() |