HEAT / app.py
Egrt's picture
修复依赖
9b51a9f
'''
Author: Egrt
Date: 2022-01-13 13:34:10
LastEditors: Egrt
LastEditTime: 2022-11-23 15:42:51
FilePath: \MaskGAN\app.py
'''
import os
os.system('cd models/ops && python3 setup.py install')
from HEAT import HEAT
import gradio as gr
heat = HEAT()
# --------模型推理---------- #
def inference(img):
image_result = heat.detect_one_image(img)
return image_result
# --------网页信息---------- #
title = "HEAT"
description = "HEAT: Holistic Edge Attention Transformer for Structured Reconstruction @Luuuu"
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.10257' target='_blank'>HEAT: Holistic Edge Attention Transformer for Structured Reconstruction </a> | <a href='https://github.com/JingyunLiang/SwinIR' target='_blank'>Github Repo</a></p>"
example_img_dir = 'images/'
example_img_name = os.listdir(example_img_dir)
examples=[[os.path.join(example_img_dir, image_path)] for image_path in example_img_name if image_path.endswith(('.jpg','.jpeg', '.png'))]
gr.Interface(
inference,
[gr.inputs.Image(type="pil", label="Input")],
gr.outputs.Image(type="pil", label="Output"),
title=title,
description=description,
article=article,
examples=examples
).launch()