Spaces:
Running
Running
| import gradio as gr | |
| import os, requests | |
| import numpy as np | |
| import torch.nn.functional as F | |
| from model.model import ResHalf | |
| from inference import Inferencer | |
| from utils import util | |
| ## local | remote | |
| RUN_MODE = "remote" | |
| if RUN_MODE != "local": | |
| os.system("wget https://huggingface.co/menghanxia/ReversibleHalftoning/resolve/main/model_best.pth.tar") | |
| os.rename("model_best.pth.tar", "./checkpoints/model_best.pth.tar") | |
| ## examples | |
| os.system("wget https://huggingface.co/menghanxia/ReversibleHalftoning/resolve/main/girl.png") | |
| os.system("wget https://huggingface.co/menghanxia/ReversibleHalftoning/resolve/main/wave.png") | |
| os.system("wget https://huggingface.co/menghanxia/ReversibleHalftoning/resolve/main/painting.png") | |
| ## step 1: set up model | |
| device = "cpu" | |
| checkpt_path = "checkpoints/model_best.pth.tar" | |
| invhalfer = Inferencer(checkpoint_path=checkpt_path, model=ResHalf(train=False), use_cuda=False, multi_gpu=False) | |
| def prepare_data(input_img, decoding_only=False): | |
| input_img = np.array(input_img / 255., np.float32) | |
| if decoding_only: | |
| input_img = input_img[:,:,:1] | |
| input_img = util.img2tensor(input_img * 2. - 1.) | |
| return input_img | |
| def run_invhalf(invhalfer, input_img, decoding_only, device="cuda"): | |
| input_img = prepare_data(input_img, decoding_only) | |
| input_img = input_img.to(device) | |
| if decoding_only: | |
| print('>>>:restoration mode') | |
| resColor = invhalfer(input_img, decoding_only=decoding_only) | |
| output = util.tensor2img(resColor / 2. + 0.5) * 255. | |
| else: | |
| print('>>>:halftoning mode') | |
| resHalftone, resColor = invhalfer(input_img, decoding_only=decoding_only) | |
| output = util.tensor2img(resHalftone / 2. + 0.5) * 255. | |
| return np.clip(output, 0, 255).astype(np.uint8) | |
| def click_run(input_img, decoding_only): | |
| output = run_invhalf(invhalfer, input_img, decoding_only, device) | |
| return output | |
| def click_move(output_img, decoding_only): | |
| if decoding_only: | |
| radio_status = "Halftoning (Photo2Halftone)" | |
| else: | |
| radio_status = "Restoration (Halftone2Photo)" | |
| return output_img, radio_status, None | |
| ## step 2: configure interface | |
| demo = gr.Blocks(title="ReversibleHalftoning") | |
| with demo: | |
| gr.Markdown(value=""" | |
| **Gradio demo for ReversibleHalftoning: Deep Halftoning with Reversible Binary Pattern**. Check our [github page](https://github.com/MenghanXia/ReversibleHalftoning) 😛. | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| Image_input = gr.Image(type="numpy", label="Input", interactive=True).style(height=480) | |
| with gr.Row(): | |
| Radio_mode = gr.Radio(type="index", choices=["Halftoning (Photo2Halftone)", "Restoration (Halftone2Photo)"], \ | |
| label="Choose a running mode", value="Halftoning (Photo2Halftone)") | |
| Button_run = gr.Button(value="Run") | |
| with gr.Column(): | |
| Image_output = gr.Image(type="numpy", label="Output").style(height=480) | |
| Button_move = gr.Button(value="Use it as input") | |
| Button_run.click(fn=click_run, inputs=[Image_input, Radio_mode], outputs=Image_output) | |
| Button_move.click(fn=click_move, inputs=[Image_output, Radio_mode], outputs=[Image_input, Radio_mode, Image_output]) | |
| if RUN_MODE != "local": | |
| gr.Examples(examples=[ | |
| ['girl.png', "Halftoning (Photo2Halftone)"], | |
| ['wave.png', "Halftoning (Photo2Halftone)"], | |
| ['painting.png', "Restoration (Halftone2Photo)"], | |
| ], | |
| inputs=[Image_input,Radio_mode], outputs=[Image_output], label="Examples") | |
| if RUN_MODE == "local": | |
| demo.launch(server_name='9.134.253.83',server_port=7788) | |
| else: | |
| demo.launch() |