|
|
import os, sys, zipfile, shutil, subprocess |
|
|
import gradio as gr |
|
|
|
|
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__)) |
|
|
ZIP_PATH = os.path.join(BASE_DIR, "engineering_drawing_code.zip") |
|
|
REPO_DIR = os.path.join(BASE_DIR, "engineering_drawing_repo") |
|
|
|
|
|
def extract_two_outputs(uploaded_image_path): |
|
|
|
|
|
if not os.path.isdir(REPO_DIR): |
|
|
with zipfile.ZipFile(ZIP_PATH, "r") as z: |
|
|
z.extractall(REPO_DIR) |
|
|
|
|
|
out_excel = os.path.join(BASE_DIR, "drawingInfo.xlsx") |
|
|
out_png = os.path.join(BASE_DIR, "extracted", "drawing01.png") |
|
|
if os.path.exists(out_excel): |
|
|
os.remove(out_excel) |
|
|
if os.path.isdir(os.path.join(BASE_DIR, "extracted")): |
|
|
shutil.rmtree(os.path.join(BASE_DIR, "extracted")) |
|
|
|
|
|
images_folder = os.path.join(REPO_DIR, "images") |
|
|
if os.path.isdir(images_folder): |
|
|
shutil.rmtree(images_folder) |
|
|
os.makedirs(images_folder, exist_ok=True) |
|
|
shutil.copyfile(uploaded_image_path, |
|
|
os.path.join(images_folder, "01.png")) |
|
|
|
|
|
subprocess.run( |
|
|
[sys.executable, "mainExtractionOCR.py"], |
|
|
cwd=REPO_DIR, |
|
|
check=True, |
|
|
) |
|
|
|
|
|
return out_png, out_excel |
|
|
|
|
|
|
|
|
demo = gr.Blocks() |
|
|
with demo: |
|
|
gr.Markdown("## Engineering-Drawing Information Extractor") |
|
|
with gr.Row(): |
|
|
inp = gr.Image(type="filepath", label="Upload drawing") |
|
|
with gr.Row(): |
|
|
balloon = gr.Image(type="filepath", label="Ballooned Drawing") |
|
|
excel = gr.File(label="Download Extracted Data (Excel)") |
|
|
btn = gr.Button("Extract") |
|
|
btn.click(fn=extract_two_outputs, inputs=inp, outputs=[balloon, excel]) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.launch() |