Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| from utils import * | |
| torch.hub.download_url_to_file( | |
| 'https://github.com/aalto-ui/aim/raw/aim2/backend/data/tests/input_values/wikipedia.org_website.png', | |
| 'wikipedia.org_website.png') | |
| torch.hub.download_url_to_file( | |
| 'https://github.com/aalto-ui/aim/raw/aim2/backend/data/tests/input_values/aalto.fi_website.png', | |
| 'aalto.fi_website.png') | |
| def inference(img, template, angel): | |
| color_image = cv2.imread(img.name, cv2.IMREAD_COLOR) | |
| height, width, _ = color_image.shape | |
| # Resize if it is bigeer than 960 * 800 | |
| if width > height: | |
| if width > 960: # 3/4 * 1280 | |
| coef_div = width / 960.0 | |
| color_image = cv2.resize(color_image, dsize=(int(width / coef_div), int(height / coef_div)), | |
| interpolation=cv2.INTER_CUBIC) | |
| else: | |
| if height > 800: # 800 | |
| coef_div = height / 800.0 | |
| color_image = cv2.resize(color_image, dsize=(int(width / coef_div), int(height / coef_div)), | |
| interpolation=cv2.INTER_CUBIC) | |
| HSV_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2HSV) | |
| selected_harmomic_scheme = HarmonicScheme(str(template), int(angel)) | |
| new_HSV_image = best_harmomic_scheme.hue_shifted(HSV_image, num_superpixels=-1) | |
| # Compute shifted histogram | |
| histo_1 = count_hue_histogram(HSV_image) | |
| histo_2 = count_hue_histogram(new_HSV_image) | |
| # Create Hue Plots | |
| fig1 = plothis(histo_1, best_harmomic_scheme, "Source Hue") | |
| fig_1_cv = get_img_from_fig(fig1) | |
| fig2 = plothis(histo_2, best_harmomic_scheme, "Target Hue") | |
| fig_2_cv = get_img_from_fig(fig2) | |
| # Stack Hue Plots | |
| vis = np.concatenate((fig_1_cv, fig_2_cv), axis=0) | |
| # Convert HSV to BGR | |
| result_image = cv2.cvtColor(new_HSV_image, cv2.COLOR_HSV2BGR) | |
| # Final output | |
| canvas = np.full((800, 960, 3), (255, 255, 255), dtype=np.uint8) | |
| # compute center offset | |
| x_center = (960 - width) // 2 | |
| y_center = (800 - height) // 2 | |
| # copy img image into center of result image | |
| canvas[y_center:y_center + height, x_center:x_center + width] = result_image | |
| # Combine | |
| output = np.concatenate((vis, canvas), axis=1) | |
| cv2.imwrite('output.png', output) | |
| return ['output.png'] | |
| title = 'Color Harmonization' | |
| description = 'Compute Color Harmonization with Different Templates' | |
| article = "<p style='text-align: center'></p>" | |
| examples = [['wikipedia.org_website.png'], ['aalto.fi_website.png']] | |
| css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}" | |
| gr.Interface( | |
| inference, | |
| [gr.inputs.Image(type='file', label='Input'), | |
| gr.inputs.Dropdown(["X", "Y", "T", "I", "mirror_L", "L", "V", "i"], | |
| default="X", | |
| label="Template"), | |
| gr.inputs.Slider(0, 359, label="Angle")], | |
| [gr.outputs.Image(type='file', label='Color Harmonization of Output Image')], | |
| title=title, | |
| description=description, | |
| article=article, | |
| examples=examples, | |
| css=css, | |
| ).launch(debug=True, enable_queue=True) | |