import os import cv2 import logging import numpy as np # --- Bắt buộc fix lỗi của HuggingFace Spaces liên quan đến việc compile Detectron2 --- # Do detectron2 yêu cầu torch lúc setup nên việc cho vào requirements.txt sẽ dính lỗi "No module torch" import subprocess try: import detectron2 except ImportError: print("Không tìm thấy detectron2, bắt đầu clone và build từ git (chỉ diễn ra 1 lần)...") subprocess.run(["pip", "install", "git+https://github.com/facebookresearch/detectron2.git", "--no-build-isolation"], check=True) import gradio as gr from huggingface_hub import hf_hub_download from pipeline_main import TechnicalDrawingPipeline # Thiết lập log level logging.basicConfig(level=logging.INFO) # Thay đổi cấu hình tải model từ mục Hugging Face Repo REPO_ID = "TVQuyet05/Detection_Tech_Draw" FILENAME = "model_final_2.pth" # Tên file weights đã được upload lên Hugging Face Hub # Cờ để kiểm tra đã load mô hình chưa pipeline = None def load_pipeline(): global pipeline if pipeline is None: print("Đang kiểm tra và tải model weights từ HuggingFace Hub...") try: # Tự động tải từ Hugging Face và lưu cache, nếu đã tải rồi nó sẽ dùng file local cache rất nhanh model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME) print(f"Đã tải thành công model tại: {model_path}") except Exception as e: raise RuntimeError(f"Lỗi tải model từ HuggingFace Hub: {e}") # Khởi tạo mô hình chỉ 1 lần duy nhất cho toàn bộ app pipeline = TechnicalDrawingPipeline(model_weights=model_path, output_dir="Pipeline_Output") return pipeline # Bảng màu hiển thị (BGR array -> RGB array) khi vẽ OpenCV COLOR_MAP = { "PartDrawing": (255, 0, 0), # Red (RGB) "Note": (0, 255, 0), # Green (RGB) "Table": (0, 0, 255) # Blue (RGB) } def draw_boxes(image_path, json_data): """ Load lại ảnh và vẽ các Bounding box minh hoạ. """ img = cv2.imread(image_path) # Chuyển BGR (OpenCV) sang RGB (Gradio) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for obj in json_data.get("objects", []): class_name = obj["class"] conf = obj["confidence"] box = obj["bbox"] x1, y1, x2, y2 = box["x1"], box["y1"], box["x2"], box["y2"] color = COLOR_MAP.get(class_name, (255, 255, 0)) # Default Yellow # Vẽ Khung cv2.rectangle(img, (x1, y1), (x2, y2), color, 3) # Vẽ Nhãn (Background mờ đằng sau nhãn cho dễ nhìn) label = f"{class_name} ({conf:.2f})" (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2) cv2.rectangle(img, (x1, y1 - 25), (x1 + w, y1), color, -1) cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) return img def extract_crops(image_path, json_data): """ Cắt ảnh ra thành các Crops và nối lại thành 1 ảnh duy nhất (Numpy Array) để hiển thị an toàn trên giao diện giống hệt cơ chế của draw_boxes. """ img = cv2.imread(image_path) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) crop_imgs = [] for obj in json_data.get("objects", []): class_name = obj["class"] conf = obj["confidence"] box = obj["bbox"] x1, y1, x2, y2 = box["x1"], box["y1"], box["x2"], box["y2"] # Đề phòng tọa độ vượt quá kích thước hoặc lỗi h, w, _ = img.shape x1, y1 = max(0, x1), max(0, y1) x2, y2 = min(w, x2), min(h, y2) if x2 > x1 and y2 > y1: c_img = img[y1:y2, x1:x2].copy() label = f"#{obj['id']} - {class_name} ({conf:.2f})" # Thêm 1 thanh header đen để ghi chữ Header bar_height = 40 c_h, c_w, _ = c_img.shape bar = np.zeros((bar_height, c_w, 3), dtype=np.uint8) cv2.putText(bar, label, (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) combined = np.vstack((bar, c_img)) crop_imgs.append(combined) if not crop_imgs: # Nếu Không có object nào return np.ones((100, 400, 3), dtype=np.uint8) * 255 # Chuẩn hoá kích thước width lớn nhất để dùng np.vstack (nối ảnh) max_w = max(c.shape[1] for c in crop_imgs) padded_crops = [] for c in crop_imgs: c_w = c.shape[1] pad_w = max_w - c_w if pad_w > 0: c = cv2.copyMakeBorder(c, 0, 0, 0, pad_w, cv2.BORDER_CONSTANT, value=[255, 255, 255]) padded_crops.append(c) # Tạo 1 bức ảnh ghép tất cả dọc từ trên xuống dưới final_img = padded_crops[0] for c in padded_crops[1:]: separator = np.ones((20, max_w, 3), dtype=np.uint8) * 150 # Gạch phân cách final_img = np.vstack((final_img, separator, c)) return final_img def format_ocr_html(json_data): """ Format kết quả OCR (text và table layout) sang định dạng HTML đẹp mắt. """ html_content = "
" objects = json_data.get("objects", []) if not objects: return "Không tìm thấy đối tượng nào." for obj in objects: cls_name = obj["class"] ocr_text = obj.get("ocr_content", "") obj_id = obj["id"] if cls_name == "Note": html_content += f"

📝 Note #{obj_id}

" if ocr_text.strip() == "": html_content += "

(Trống)

" else: html_content += f"
{ocr_text}
" elif cls_name == "Table": html_content += f"

📊 Table #{obj_id} (HTML)

" if ocr_text.strip() == "": html_content += "

(Không giải mã được cấu trúc HTML)

" else: # Add table border styling inline for better preview styled_table = ocr_text.replace("", "").replace("", "") styled_table = styled_table.replace("{str(e)}" # Gọi inference chính results = pipe.process_image(image_path) if not results: return None, {}, None, "Có lỗi khi phân tích ảnh này." # Tạo kết quả thị giác hóa (Visualization) annotated_img = draw_boxes(image_path, results) # Cắt ảnh crop object & Gộp thành 1 ảnh numpy array duy nhất ! collage_crops_img = extract_crops(image_path, results) # Tạo HTML render OCR ocr_html = format_ocr_html(results) return annotated_img, results, collage_crops_img, ocr_html # Xây dựng giao diện web với Gradio with gr.Blocks(title="Technical Draw Detection") as demo: gr.Markdown("# 📐 Technical Drawing Layout Analysis") gr.Markdown("Nhận diện bố cục (PartDrawing, Note, Table) trên bản vẽ kỹ thuật, trích xuất text (PP-OCR) và mô hình hóa Bảng (SLANet). **Optimized for CPU.**") with gr.Row(): # Cột Input with gr.Column(scale=4): input_image = gr.Image(type="filepath", label="Upload ảnh bản vẽ") submit_btn = gr.Button("🚀 Gửi / Detect", variant="primary") # Cột Output (Visual & JSON) with gr.Column(scale=6): output_image = gr.Image(label="Visualization") json_output = gr.JSON(label="JSON Panel (Thông tin toạ độ & Meta)") # Hàng ở dưới full length cho Crops (Dùng chung tính năng gr.Image) with gr.Row(): with gr.Column(scale=10): gr.Markdown("## 🖼️ Các đối tượng được tách ra (Cropped Objects)") crop_collage_ui = gr.Image(label="Cropped Objects Visualization") # Hàng ở dưới full length cho OCR with gr.Row(): with gr.Column(scale=10): gr.Markdown("## 🖹 OCR & Table HTML Panel") ocr_output = gr.HTML(label="OCR results") # Gán sự kiện submit_btn.click( fn=process_ui, inputs=[input_image], outputs=[output_image, json_output, crop_collage_ui, ocr_output] ) if __name__ == "__main__": # Launch server # Nếu deploy HuggingSpaces thì để mặc định 0.0.0.0 demo.launch(server_name="0.0.0.0", server_port=7860, share=False)