Refactor code

CRAFT/__init__.py

@@ -1,2 +1,2 @@
-from .model import CRAFTModel
+from .model import CRAFTModel, preprocess_image
 from .utils import draw_boxes, draw_polygons, boxes_area, polygons_area

CRAFT/model.py

@@ -122,9 +122,7 @@ class CRAFTModel:
             res.append(poly.astype(np.int32).tolist())
         return res
     
-    def get_boxes(self, image: Image.Image) -> List[List[List[int]]]:
-        x, ratio_w, ratio_h = preprocess_image(np.array(image), self.canvas_size, self.mag_ratio)
-        
+    def _get_boxes_preproc(self, x, ratio_w, ratio_h) -> List[List[List[int]]]:
         score_text, score_link = self.get_text_map(x, ratio_w, ratio_h)
         
         # Post-processing
@@ -141,4 +139,10 @@ class CRAFTModel:
             for box in boxes:
                 boxes_final.append([box[0], box[2]])
 
-        return boxes_final
+        return boxes_final
+    
+    def get_boxes(self, image: Image.Image) -> List[List[List[int]]]:
+        x, ratio_w, ratio_h = preprocess_image(np.array(image), self.canvas_size, self.mag_ratio)
+        
+        boxes_final = self._get_boxes_preproc(x, ratio_w, ratio_h)
+        return boxes_final

CRAFT/utils.py

@@ -48,21 +48,19 @@ def box_to_poly(box: List[List[int]]) -> List[List[int]]:
     return [box[0], [box[0][0], box[1][1]], box[1], [box[1][0], box[0][1]]]
 
 
-def boxes_area(img: Image.Image, bboxes: List[List[List[int]]]) -> int:
-    img_s = img.size[0]*img.size[1]
+def boxes_area(bboxes: List[List[List[int]]]) -> int:
     total_S = 0
     for box in bboxes:
         pgon = Polygon(box_to_poly(box)) 
         S = pgon.area
         total_S+=S
-    return total_S/img_s
+    return total_S
 
 
-def polygons_area(img: Image.Image, polygons: List[List[List[int]]]) -> int:
-    img_s = img.size[0]*img.size[1]
+def polygons_area(polygons: List[List[List[int]]]) -> int:
     total_S = 0
     for poly in polygons:
         pgon = Polygon(poly) 
         S = pgon.area
         total_S+=S
-    return total_S/img_s
+    return total_S

example.ipynb

@@ -3,7 +3,7 @@
   {
    "cell_type": "code",
    "execution_count": 1,
-   "id": "2cdb2ec0",
+   "id": "bba67631",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -16,7 +16,7 @@
   {
    "cell_type": "code",
    "execution_count": 2,
-   "id": "7860eda1",
+   "id": "6045288a",
    "metadata": {},
    "outputs": [
     {
@@ -35,7 +35,7 @@
   {
    "cell_type": "code",
    "execution_count": 3,
-   "id": "f710c8e7",
+   "id": "000251de",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -44,7 +44,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "704faeea",
+   "id": "1f6755fb",
    "metadata": {},
    "source": [
     "## Predict boxes"
@@ -52,16 +52,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
-   "id": "07a7d242",
+   "execution_count": 14,
+   "id": "a00bd58a",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "CPU times: user 62.3 ms, sys: 31.4 ms, total: 93.7 ms\n",
-      "Wall time: 61.7 ms\n"
+      "CPU times: user 77.6 ms, sys: 20.9 ms, total: 98.5 ms\n",
+      "Wall time: 67.8 ms\n"
      ]
     }
    ],
@@ -74,7 +74,7 @@
   {
    "cell_type": "code",
    "execution_count": 5,
-   "id": "755611db",
+   "id": "10c9ffcf",
    "metadata": {},
    "outputs": [
     {
@@ -96,7 +96,7 @@
   {
    "cell_type": "code",
    "execution_count": 6,
-   "id": "e7303093",
+   "id": "b69f82fd",
    "metadata": {},
    "outputs": [
     {
@@ -113,12 +113,12 @@
    "source": [
     "# calculating text area in percentage to total image area\n",
     "\n",
-    "boxes_area(img, boxes)"
+    "boxes_area(boxes)/(img.size[0]*img.size[1])"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "dcba5286",
+   "id": "1b9e564c",
    "metadata": {},
    "source": [
     "## Predict polygons"
@@ -126,16 +126,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
-   "id": "d802595d",
+   "execution_count": 8,
+   "id": "65751da8",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "CPU times: user 76.2 ms, sys: 54.6 ms, total: 131 ms\n",
-      "Wall time: 75.8 ms\n"
+      "CPU times: user 77 ms, sys: 57.7 ms, total: 135 ms\n",
+      "Wall time: 80.3 ms\n"
      ]
     },
     {
@@ -155,8 +155,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
-   "id": "784fa82a",
+   "execution_count": 9,
+   "id": "f18ababd",
    "metadata": {},
    "outputs": [
     {
@@ -166,7 +166,7 @@
        "<PIL.Image.Image image mode=RGB size=500x603>"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 9,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -177,8 +177,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
-   "id": "635682fd",
+   "execution_count": 10,
+   "id": "183a32f0",
    "metadata": {},
    "outputs": [
     {
@@ -187,7 +187,7 @@
        "0.23886567164179104"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -195,13 +195,13 @@
    "source": [
     "# calculating text area in percentage to total image area\n",
     "\n",
-    "polygons_area(img, polygons)"
+    "polygons_area(polygons)/(img.size[0]*img.size[1])"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0db11b08",
+   "id": "17190321",
    "metadata": {},
    "outputs": [],
    "source": []