cherubicxn/ScaleLSD
2
1import torch2import cv23import os4import gradio as gr5import numpy as np6import random7from pathlib import Path8import json9import spaces10 11 12# Title for the Gradio interface13_TITLE = 'Gradio Demo of ScaleLSD for Structured Representation of Images'14MAX_SEED = 100015 16os.system('mkdir -p models')17os.system('wget https://huggingface.co/cherubicxn/scalelsd/resolve/main/scalelsd-vitbase-v2-train-sa1b.pt -O models/scalelsd-vitbase-v2-train-sa1b.pt')18os.system('wget https://huggingface.co/cherubicxn/scalelsd/resolve/main/scalelsd-vitbase-v1-train-sa1b.pt -O models/scalelsd-vitbase-v1-train-sa1b.pt')19os.system('pip install -e .')20 21 22def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:23 """random seed"""24 if randomize_seed:25 seed = random.randint(0, MAX_SEED)26 return seed27 28def stop_run():29 """stop run"""30 return (31 gr.update(value="Run", variant="primary", visible=True),32 gr.update(visible=False),33 )34 35# @spaces.GPU()36@spaces.GPU37def process_image(38 input_image,39 model_name='scalelsd-vitbase-v2-train-sa1b.pt',40 save_name='temp_output',41 threshold=10,42 junction_threshold_hm=0.008,43 num_junctions_inference=512,44 width=512,45 height=512,46 line_width=2,47 juncs_size=4,48 whitebg=0.0,49 draw_junctions_only=False,50 use_lsd=False,51 use_nms=False,52 edge_color='orange',53 vertex_color='Cyan',54 output_format='png',55 seed=0,56 randomize_seed=False57):58 use_lsd = False59 from scalelsd.ssl.models.detector import ScaleLSD60 from scalelsd.base import show, WireframeGraph61 from scalelsd.ssl.misc.train_utils import fix_seeds, load_scalelsd_model62 """core processing function for image inference"""63 # set random seed64 seed = int(randomize_seed_fn(seed, randomize_seed))65 fix_seeds(seed)66 67 # initialize model68 ckpt = "models/" + model_name69 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")70 model = load_scalelsd_model(ckpt, device)71 72 # set model parameters73 model.junction_threshold_hm = junction_threshold_hm74 model.num_junctions_inference = num_junctions_inference75 76 # transform input image77 if isinstance(input_image, np.ndarray):78 image = cv2.cvtColor(input_image, cv2.COLOR_RGB2GRAY)79 else:80 image = cv2.imread(input_image, 0)81 82 # resize83 ori_shape = image.shape[:2]84 image_resized = cv2.resize(image.copy(), (width, height))85 image_tensor = torch.from_numpy(image_resized).float() / 255.086 image_tensor = image_tensor[None, None].to('cuda')87 88 # meta data89 meta = {90 'width': ori_shape[1],91 'height': ori_shape[0],92 'filename': '',93 'use_lsd': use_lsd,94 'use_nms': use_nms,95 }96 97 # inference98 with torch.no_grad():99 outputs, _ = model(image_tensor, meta)100 outputs = outputs[0]101 102 # visual results103 painter = show.painters.HAWPainter()104 painter.confidence_threshold = threshold105 painter.line_width = line_width106 painter.marker_size = juncs_size107 if whitebg > 0.0:108 show.Canvas.white_overlay = whitebg109 110 temp_folder = "temp_output"111 os.makedirs(temp_folder, exist_ok=True)112 fig_file = f"{temp_folder}/{save_name}.png"113 with show.image_canvas(input_image, fig_file=fig_file) as ax:114 if draw_junctions_only:115 painter.draw_junctions(ax, outputs)116 else:117 painter.draw_wireframe(ax, outputs, edge_color=edge_color, vertex_color=vertex_color)118 # read the result image119 result_image = cv2.imread(fig_file)120 121 if output_format != 'png':122 fig_file = f"{temp_folder}/{save_name}.{output_format}"123 with show.image_canvas(input_image, fig_file=fig_file) as ax:124 if draw_junctions_only:125 painter.draw_junctions(ax, outputs)126 else:127 painter.draw_wireframe(ax, outputs, edge_color=edge_color, vertex_color=vertex_color)128 129 json_file = f"{temp_folder}/{save_name}.json"130 indices = WireframeGraph.xyxy2indices(outputs['juncs_pred'],outputs['lines_pred'])131 wireframe = WireframeGraph(outputs['juncs_pred'], outputs['juncs_score'], indices, outputs['lines_score'], outputs['width'], outputs['height'])132 with open(json_file, 'w') as f:133 json.dump(wireframe.jsonize(),f)134 135 136 return result_image[:, :, ::-1], json_file, fig_file137 138def run_demo():139 """create the Gradio demo interface"""140 css = """141 #col-container {142 margin: 0 auto;143 max-width: 800px;144 }145 """146 147 with gr.Blocks(css=css, title=_TITLE) as demo:148 with gr.Column(elem_id="col-container"):149 gr.Markdown(f'# {_TITLE}')150 gr.Markdown("Detect wireframe structures in images using ScaleLSD model")151 152 pid = gr.State()153 figs_root = "assets/figs"154 example_images = [os.path.join(figs_root, iname) for iname in os.listdir(figs_root)]155 156 with gr.Row():157 input_image = gr.Image(example_images[0], label="Input Image", type="numpy")158 output_image = gr.Image(label="Detection Result")159 160 with gr.Row():161 run_btn = gr.Button(value="Run", variant="primary")162 stop_btn = gr.Button(value="Stop", variant="stop", visible=False)163 164 with gr.Row():165 json_file = gr.File(label="Download JSON Output", type="filepath")166 image_file = gr.File(label="Download Image Output", type="filepath")167 168 with gr.Accordion("Advanced Settings", open=True):169 with gr.Row():170 model_name = gr.Dropdown(171 [ckpt for ckpt in os.listdir('models') if ckpt.endswith('.pt')],172 value='scalelsd-vitbase-v2-train-sa1b.pt', 173 label="Model Selection"174 )175 176 with gr.Row():177 save_name = gr.Textbox('temp_output', label="Save Name", placeholder="Name for saving output files")178 179 with gr.Row():180 with gr.Column():181 threshold = gr.Number(10, label="Line Threshold")182 junction_threshold_hm = gr.Number(0.008, label="Junction Threshold")183 num_junctions_inference = gr.Number(1024, label="Max Number of Junctions")184 width = gr.Number(512, label="Input Width")185 height = gr.Number(512, label="Input Height")186 187 with gr.Column():188 draw_junctions_only = gr.Checkbox(False, label="Show Junctions Only")189 use_lsd = gr.Checkbox(False, label="Use LSD-Rectifier")190 use_nms = gr.Checkbox(True, label="Use NMS")191 output_format = gr.Dropdown(192 ['png', 'jpg', 'pdf'], 193 value='png', 194 label="Output Format"195 )196 whitebg = gr.Slider(0.0, 1.0, value=0.7, label="White Background Opacity")197 line_width = gr.Number(2, label="Line Width")198 juncs_size = gr.Number(8, label="Junctions Size")199 200 with gr.Row():201 edge_color = gr.Dropdown(202 ['orange', 'midnightblue', 'red', 'green'], 203 value='orange', 204 label="Edge Color"205 )206 vertex_color = gr.Dropdown(207 ['Cyan', 'deeppink', 'yellow', 'purple'], 208 value='Cyan', 209 label="Vertex Color"210 )211 212 with gr.Row():213 randomize_seed = gr.Checkbox(False, label="Randomize Seed")214 seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")215 216 gr.Examples(217 examples=example_images,218 inputs=input_image,219 )220 221 # star event handlers222 run_event = run_btn.click(223 fn=process_image,224 inputs=[225 input_image,226 model_name,227 save_name,228 threshold,229 junction_threshold_hm,230 num_junctions_inference,231 width,232 height,233 line_width,234 juncs_size,235 whitebg,236 draw_junctions_only,237 use_lsd,238 use_nms,239 edge_color,240 vertex_color,241 output_format,242 seed,243 randomize_seed244 ],245 outputs=[output_image, json_file, image_file],246 )247 248 # stop event handlers249 stop_btn.click(250 fn=stop_run,251 outputs=[run_btn, stop_btn],252 cancels=[run_event],253 queue=False,254 )255 256 257 return demo258 259run_demo().launch()260