Yolo / app.py
ngohel58's picture
Update app.py
1db73f8 verified
raw
history blame
1.84 kB
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):
# 1) Unzip the extractor repo if needed
if not os.path.isdir(REPO_DIR):
with zipfile.ZipFile(ZIP_PATH, "r") as z:
z.extractall(REPO_DIR)
# 2) Clean previous outputs
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"))
# 3) Prepare images folder
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"))
# 4) Run the OCR/extraction script
subprocess.run(
[sys.executable, "mainExtractionOCR.py"],
cwd=REPO_DIR,
check=True,
)
# 5) Return the ballooned drawing + Excel
return out_png, out_excel
# Build Gradio UI
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()