TAK

Drawings_RT_DETR_Demo.ipynb

@@ -0,0 +1,201 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# Bản Demo Nhận diện Bản vẽ Kỹ thuật (RT-DETR + EasyOCR)\n",
+        "Sử dụng Google Colab để chạy thử mô hình của bạn một cách nhanh chóng."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "!pip install ultralytics easyocr gdown opencv-python-headless matplotlib"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import gdown\n",
+        "import os\n",
+        "\n",
+        "# Thay bằng ID file best.pt trên Google Drive của bạn\n",
+        "file_id = '1WnIVT9nI7uDmfKzsWw8aTFZQ1VAuVF06'\n",
+        "url = f'https://drive.google.com/uc?id={file_id}'\n",
+        "\n",
+        "output = 'best.pt'\n",
+        "if not os.path.exists(output):\n",
+        "    print(\"Đang tải model weights...\")\n",
+        "    gdown.download(url, output, quiet=False)\n",
+        "    print(\"Đã tải xong!\")\n",
+        "else:\n",
+        "    print(\"Model đã tồn tại.\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import cv2\n",
+        "import easyocr\n",
+        "import matplotlib.pyplot as plt\n",
+        "from ultralytics import RTDETR\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "\n",
+        "print(\"Loading RT-DETR...\")\n",
+        "model = RTDETR('best.pt')\n",
+        "\n",
+        "print(\"Loading EasyOCR...\")\n",
+        "reader = easyocr.Reader(['vi', 'en'], gpu=True) # Trên Colab nhớ chọn Runtime -> GPU\n",
+        "print(\"Ready!\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Chạy cell này để upload bức ảnh bản vẽ của bạn lên Colab\n",
+        "uploaded = files.upload()\n",
+        "image_path = list(uploaded.keys())[0]\n",
+        "print(f\"Đã tải lên ảnh: {image_path}\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "def perform_easyocr_on_crop(cropped_img, is_table=False):\n",
+        "    if cropped_img is None or cropped_img.size == 0:\n",
+        "        return \"\"\n",
+        "    try:\n",
+        "        gray_img = cv2.cvtColor(cropped_img, cv2.COLOR_BGR2GRAY)\n",
+        "        blurred_img = cv2.GaussianBlur(gray_img, (3, 3), 0)\n",
+        "        _, thresh_img = cv2.threshold(blurred_img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)\n",
+        "\n",
+        "        if not is_table:\n",
+        "            result = reader.readtext(thresh_img, detail=0, paragraph=True)\n",
+        "            return \" \".join(result) if result else \"\"\n",
+        "        else:\n",
+        "            result = reader.readtext(thresh_img, detail=1)\n",
+        "            if not result:\n",
+        "                return \"\"\n",
+        "            lines = []\n",
+        "            for bbox, text, conf in result:\n",
+        "                y_center = (bbox[0][1] + bbox[2][1]) / 2\n",
+        "                x_center = (bbox[0][0] + bbox[1][0]) / 2\n",
+        "                lines.append({\"y\": y_center, \"x\": x_center, \"text\": text})\n",
+        "            lines.sort(key=lambda item: item['y'])\n",
+        "            y_tolerance = 20\n",
+        "            rows = []\n",
+        "            current_row = []\n",
+        "            current_y = None\n",
+        "            for item in lines:\n",
+        "                if current_y is None:\n",
+        "                    current_y = item['y']\n",
+        "                    current_row.append(item)\n",
+        "                elif abs(item['y'] - current_y) <= y_tolerance:\n",
+        "                    current_row.append(item)\n",
+        "                    current_y = (current_y * (len(current_row) - 1) + item['y']) / len(current_row)\n",
+        "                else:\n",
+        "                    rows.append(current_row)\n",
+        "                    current_row = [item]\n",
+        "                    current_y = item['y']\n",
+        "            if current_row:\n",
+        "                rows.append(current_row)\n",
+        "            formatted_text = \"\"\n",
+        "            for row in rows:\n",
+        "                row.sort(key=lambda item: item['x'])\n",
+        "                row_text = \" | \".join([item['text'] for item in row])\n",
+        "                formatted_text += row_text + \"\\n\"\n",
+        "            return formatted_text.strip()\n",
+        "    except Exception as e:\n",
+        "        print(f\"Lỗi OCR: {e}\")\n",
+        "        return \"\"\n",
+        "\n",
+        "# Đọc ảnh\n",
+        "cv_img = cv2.imread(image_path)\n",
+        "\n",
+        "# Dự đoán\n",
+        "results = model.predict(image_path, conf=0.25, iou=0.45, verbose=False)\n",
+        "boxes = results[0].boxes\n",
+        "\n",
+        "print(\"\\n--- KẾT QUẢ NHẬN DIỆN ---\")\n",
+        "if boxes is not None and len(boxes) > 0:\n",
+        "    for box in boxes:\n",
+        "        x1, y1, x2, y2 = box.xyxy[0].tolist()\n",
+        "        conf_val = float(box.conf[0])\n",
+        "        cls_id = int(box.cls[0])\n",
+        "        cls_name = model.names[cls_id]\n",
+        "        cls_lower = cls_name.lower()\n",
+        "\n",
+        "        if cls_lower == 'table':\n",
+        "            color = (0, 0, 255) # Đỏ\n",
+        "            thickness = 3\n",
+        "        elif cls_lower == 'note':\n",
+        "            color = (0, 150, 0) # Xanh lá đậm\n",
+        "            thickness = 2\n",
+        "        elif cls_lower == 'partdrawing':\n",
+        "            color = (255, 0, 0) # Xanh biển\n",
+        "            thickness = 2\n",
+        "        else:\n",
+        "            color = (0, 165, 255) # Cam\n",
+        "            thickness = 2\n",
+        "\n",
+        "        # Crop\n",
+        "        x1i, y1i, x2i, y2i = max(0, int(x1)), max(0, int(y1)), min(cv_img.shape[1], int(x2)), min(cv_img.shape[0], int(y2))\n",
+        "        cropped_img = cv_img[y1i:y2i, x1i:x2i]\n",
+        "        \n",
+        "        ocr_text = \"\"\n",
+        "        if cls_lower in ['table', 'note'] and cropped_img.size != 0:\n",
+        "            ocr_text = perform_easyocr_on_crop(cropped_img, is_table=(cls_lower == 'table'))\n",
+        "\n",
+        "        print(f\"- {cls_name} (Độ tin cậy: {conf_val:.2f})\")\n",
+        "        if ocr_text:\n",
+        "            print(f\"  [OCR] Nội dung:\\n{ocr_text}\\n\")\n",
+        "\n",
+        "        # Vẽ khung trên ảnh\n",
+        "        cv2.rectangle(cv_img, (x1i, y1i), (x2i, y2i), color, thickness)\n",
+        "        label = f\"{cls_name} {conf_val:.2f}\"\n",
+        "        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)\n",
+        "        label_y = max(y1i, 20)\n",
+        "        cv2.rectangle(cv_img, (x1i, label_y - 20), (x1i + tw, label_y), color, -1)\n",
+        "        cv2.putText(cv_img, label, (x1i, label_y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)\n",
+        "\n",
+        "# Hiển thị ảnh kết quả\n",
+        "plt.figure(figsize=(15, 15))\n",
+        "plt.imshow(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))\n",
+        "plt.axis('off')\n",
+        "plt.show()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "Drawings_RT_DETR_Demo.ipynb",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

app.py

@@ -137,7 +137,7 @@ def predict():
 
         # Run inference
         t0 = time.time()
-        results = model.predict(image, conf=conf, iou=iou, imgsz=640, verbose=False)
+        results = model.predict(image, conf=conf, iou=iou, verbose=False)
         inference_time = round((time.time() - t0) * 1000, 1)
 
         result = results[0]
@@ -154,18 +154,17 @@ def predict():
                 ocr_text = ""
 
                 # --- Xử lý vẽ khung ---
-                cls_lower = cls_name.lower()
-                if cls_lower == 'table':
-                    color = (0, 0, 255)    # Đỏ (BGR) cho Table
+                if cls_name == 'Table':
+                    color = (0, 0, 255)    
                     thickness = 3          
-                elif cls_lower == 'note':
-                    color = (0, 150, 0)    # Xanh lá đậm (BGR) cho Note
+                elif cls_name == 'Note':
+                    color = (0, 255, 0)
                     thickness = 2
-                elif cls_lower == 'partdrawing':
-                    color = (255, 0, 0)    # Xanh biển (BGR) cho PartDrawing
+                elif cls_name == 'PartDrawing':
+                    color = (255, 0, 0)    
                     thickness = 2
                 else:
-                    color = (0, 165, 255)  # Cam (BGR) cho class khác
+                    color = (0, 255, 255)  
                     thickness = 2
 
                 crop_filename = ""
@@ -186,8 +185,8 @@ def predict():
                         crop_b64 = base64.b64encode(buffer).decode('utf-8')
 
                         # --- GỌI EASYOCR ĐỂ ĐỌC CHỮ (CHỈ CHO NOTE VÀ TABLE) ---
-                        if cls_lower in ['table', 'note']:
-                            ocr_text = perform_easyocr_on_crop(cropped_img, is_table=(cls_lower == 'table'))
+                        if cls_name in ['Table', 'Note']:
+                            ocr_text = perform_easyocr_on_crop(cropped_img, is_table=(cls_name == 'Table'))
                             if ocr_text:
                                 print(f"[OCR SUCCESS] Đã đọc được đoạn: {len(ocr_text)} characters")
                     else:

test_predict.py

@@ -0,0 +1,22 @@
+import requests
+import sys
+
+try:
+    with open("uploads/test.jpg", "wb") as f:
+        f.write(b"fake image data")
+except Exception:
+    pass
+
+# We need a proper image to test the API. Let's create a blank image using numpy and cv2
+import cv2
+import numpy as np
+blank_image = np.zeros((500, 500, 3), np.uint8)
+cv2.imwrite("blank_test.jpg", blank_image)
+
+url = "http://127.0.0.1:5002/predict"
+files = {'image': open('blank_test.jpg', 'rb')}
+data = {'conf': 0.1, 'iou': 0.1}
+
+response = requests.post(url, files=files, data=data)
+print("Status Code:", response.status_code)
+print("Response:", response.text[:500])