Spaces:

phamha
/

engineering-drawing-analyzer

Configuration error

Harry Pham

init project

ea9cf0f 7 months ago

4.08 kB

# app.py  ← đặt ở root project
import gradio as gr
import json, shutil, tempfile
from pathlib import Path
from PIL import Image
import cv2, numpy as np

from src.inference import run_pipelineas

# Auto-download weights nếu chưa có
import os
if not os.path.exists("best.pt"):
    from huggingface_hub import hf_hub_download
    hf_hub_download(
        repo_id="phamha/drawing-model-weights",
        filename="best.pt",
        local_dir="."
    )

CHECKPOINT = "best.pt"

# ── Hàm xử lý cho Gradio ─────────────────────────────────
def process(image: Image.Image):
    if image is None:
        return None, "{}", "Chưa có ảnh."

    # Lưu ảnh PIL tạm
    tmp_dir = tempfile.mkdtemp()
    tmp_path = f"{tmp_dir}/input.jpg"
    image.save(tmp_path, quality=95)

    try:
        result, vis_path = run_pipeline(
            image_path=tmp_path,
            output_dir=tmp_dir,
            checkpoint=CHECKPOINT,
            conf=0.3,
        )
    except Exception as e:
        return None, "{}", f"Lỗi: {e}"

    # ── Ảnh visualize ─────────────────────────────────────
    vis_bgr = cv2.imread(vis_path)
    vis_rgb = cv2.cvtColor(vis_bgr, cv2.COLOR_BGR2RGB)

    # ── JSON (bỏ crop_path để gọn) ───────────────────────
    clean = {"image": result["image"], "objects": []}
    for obj in result["objects"]:
        clean["objects"].append({
            "id":          obj["id"],
            "class":       obj["class"],
            "confidence":  obj["confidence"],
            "bbox":        obj["bbox"],
            "ocr_content": obj["ocr_content"],
        })
    json_str = json.dumps(clean, ensure_ascii=False, indent=2)

    # ── OCR text đẹp ──────────────────────────────────────
    ocr_parts = []
    for obj in result["objects"]:
        if not obj["ocr_content"]:
            continue
        content = obj["ocr_content"]
        if isinstance(content, dict):       # Table
            content = content.get("text", "")
        sep = "─" * 44
        ocr_parts.append(
            f"{sep}\n"
            f"[{obj['class']} #{obj['id']}]  "
            f"conf={obj['confidence']}\n"
            f"{sep}\n{content}"
        )
    ocr_text = "\n\n".join(ocr_parts) if ocr_parts else "Không có vùng Note/Table."

    return vis_rgb, json_str, ocr_text


# ── Giao diện Gradio ──────────────────────────────────────
with gr.Blocks(title="Engineering Drawing Analyzer") as demo:
    gr.Markdown("""
    # 🔧 Engineering Drawing Analyzer
    Tự động phát hiện và trích xuất **PartDrawing · Note · Table** từ bản vẽ kỹ thuật.
    """)

    with gr.Row():
        with gr.Column(scale=1):
            inp_image = gr.Image(
                type="pil",
                label="Upload bản vẽ kỹ thuật",
            )
            btn = gr.Button("🔍 Detect & OCR", variant="primary")

        with gr.Column(scale=1):
            out_image = gr.Image(label="Kết quả detection")

    with gr.Row():
        with gr.Column(scale=1):
            out_json = gr.Code(
                language="json",
                label="JSON output",
                lines=20,
            )
        with gr.Column(scale=1):
            out_ocr = gr.Textbox(
                label="OCR content (Note & Table)",
                lines=20,
                max_lines=40,
            )

    btn.click(
        fn=process,
        inputs=[inp_image],
        outputs=[out_image, out_json, out_ocr],
    )

    gr.Examples(
        examples=[["test.jpg"]],    # thêm ảnh mẫu nếu có
        inputs=[inp_image],
    )

if __name__ == "__main__":
    demo.launch(
        share=True,     # tạo link public tạm thời
        server_port=7860,
    )

	# app.py ← đặt ở root project
	import gradio as gr
	import json, shutil, tempfile
	from pathlib import Path
	from PIL import Image
	import cv2, numpy as np

	from src.inference import run_pipelineas

	# Auto-download weights nếu chưa có
	import os
	if not os.path.exists("best.pt"):
	from huggingface_hub import hf_hub_download
	hf_hub_download(
	repo_id="phamha/drawing-model-weights",
	filename="best.pt",
	local_dir="."
	)

	CHECKPOINT = "best.pt"

	# ── Hàm xử lý cho Gradio ─────────────────────────────────
	def process(image: Image.Image):
	if image is None:
	return None, "{}", "Chưa có ảnh."

	# Lưu ảnh PIL tạm
	tmp_dir = tempfile.mkdtemp()
	tmp_path = f"{tmp_dir}/input.jpg"
	image.save(tmp_path, quality=95)

	try:
	result, vis_path = run_pipeline(
	image_path=tmp_path,
	output_dir=tmp_dir,
	checkpoint=CHECKPOINT,
	conf=0.3,
	)
	except Exception as e:
	return None, "{}", f"Lỗi: {e}"

	# ── Ảnh visualize ─────────────────────────────────────
	vis_bgr = cv2.imread(vis_path)
	vis_rgb = cv2.cvtColor(vis_bgr, cv2.COLOR_BGR2RGB)

	# ── JSON (bỏ crop_path để gọn) ───────────────────────
	clean = {"image": result["image"], "objects": []}
	for obj in result["objects"]:
	clean["objects"].append({
	"id": obj["id"],
	"class": obj["class"],
	"confidence": obj["confidence"],
	"bbox": obj["bbox"],
	"ocr_content": obj["ocr_content"],
	})
	json_str = json.dumps(clean, ensure_ascii=False, indent=2)

	# ── OCR text đẹp ──────────────────────────────────────
	ocr_parts = []
	for obj in result["objects"]:
	if not obj["ocr_content"]:
	continue
	content = obj["ocr_content"]
	if isinstance(content, dict): # Table
	content = content.get("text", "")
	sep = "─" * 44
	ocr_parts.append(
	f"{sep}\n"
	f"[{obj['class']} #{obj['id']}] "
	f"conf={obj['confidence']}\n"
	f"{sep}\n{content}"
	)
	ocr_text = "\n\n".join(ocr_parts) if ocr_parts else "Không có vùng Note/Table."

	return vis_rgb, json_str, ocr_text


	# ── Giao diện Gradio ──────────────────────────────────────
	with gr.Blocks(title="Engineering Drawing Analyzer") as demo:
	gr.Markdown("""
	# 🔧 Engineering Drawing Analyzer
	Tự động phát hiện và trích xuất PartDrawing · Note · Table từ bản vẽ kỹ thuật.
	""")

	with gr.Row():
	with gr.Column(scale=1):
	inp_image = gr.Image(
	type="pil",
	label="Upload bản vẽ kỹ thuật",
	)
	btn = gr.Button("🔍 Detect & OCR", variant="primary")

	with gr.Column(scale=1):
	out_image = gr.Image(label="Kết quả detection")

	with gr.Row():
	with gr.Column(scale=1):
	out_json = gr.Code(
	language="json",
	label="JSON output",
	lines=20,
	)
	with gr.Column(scale=1):
	out_ocr = gr.Textbox(
	label="OCR content (Note & Table)",
	lines=20,
	max_lines=40,
	)

	btn.click(
	fn=process,
	inputs=[inp_image],
	outputs=[out_image, out_json, out_ocr],
	)

	gr.Examples(
	examples=[["test.jpg"]], # thêm ảnh mẫu nếu có
	inputs=[inp_image],
	)

	if __name__ == "__main__":
	demo.launch(
	share=True, # tạo link public tạm thời
	server_port=7860,
	)