From 9f55a0c034a2fe14ceb3ab9a6260ac625421fdb9 Mon Sep 17 00:00:00 2001
From: Gosia <malrog1@st.amu.edu.pl>
Date: Tue, 19 Dec 2023 09:32:05 +0100
Subject: [PATCH] BLUR-37: Refactor code to move creation of image with boxes
 to graphics directory.

---
 graphics/add_boxes_to_image.py | 26 ++++++++++++++++++++++++++
 ml/element_detection.py        | 16 ----------------
 ml/test.py                     |  4 +++-
 3 files changed, 29 insertions(+), 17 deletions(-)
 create mode 100644 graphics/add_boxes_to_image.py

diff --git a/graphics/add_boxes_to_image.py b/graphics/add_boxes_to_image.py
new file mode 100644
index 0000000..e8b3eec
--- /dev/null
+++ b/graphics/add_boxes_to_image.py
@@ -0,0 +1,26 @@
+import os
+
+from typing import List
+
+from PIL import Image, ImageDraw, ImageFont
+
+from BlurMe.ml.element_detection import BoundBox
+
+DIR_PATH = os.path.dirname(os.path.realpath(__file__))
+
+
+def show_image_with_boxes(
+    in_image_path: str, bounding_boxes: List[BoundBox], out_image_path: str = None
+):
+    img = Image.open(in_image_path)
+    draw = ImageDraw.Draw(img)
+    font_path = DIR_PATH + "/assets/fonts/arial.ttf"
+    font = ImageFont.truetype(font_path, 25)
+    for i, box in enumerate(bounding_boxes):
+        draw.rectangle(box.get_params(), outline="red", width=2, fill=None)
+        draw.text((box.x1 + 5, box.y1 + 5), str(i+1), fill="red", font=font)
+    if not out_image_path:
+        out_image_path = (
+            in_image_path.split(".")[0] + "_out." + in_image_path.split(".")[1]
+        )
+    img.save(out_image_path)
diff --git a/ml/element_detection.py b/ml/element_detection.py
index 020f561..6cfd74d 100644
--- a/ml/element_detection.py
+++ b/ml/element_detection.py
@@ -2,7 +2,6 @@ import os
 from typing import List, Tuple
 
 import torch
-from PIL import Image, ImageDraw, ImageFont
 
 from ultralytics import YOLO
 
@@ -55,18 +54,3 @@ def detect(image_path: str) -> List[BoundBox]:
     return bounding_boxes
 
 
-def show_image_with_boxes(
-    in_image_path: str, bounding_boxes: List[BoundBox], out_image_path: str = None
-):
-    img = Image.open(in_image_path)
-    draw = ImageDraw.Draw(img)
-    font_path = DIR_PATH + "/assets/fonts/arial.ttf"
-    font = ImageFont.truetype(font_path, 25)
-    for i, box in enumerate(bounding_boxes):
-        draw.rectangle(box.get_params(), outline="red", width=2, fill=None)
-        draw.text((box.x1 + 5, box.y1 + 5), str(i+1), fill="red", font=font)
-    if not out_image_path:
-        out_image_path = (
-            in_image_path.split(".")[0] + "_out." + in_image_path.split(".")[1]
-        )
-    img.save(out_image_path)
diff --git a/ml/test.py b/ml/test.py
index e41296a..ad99d45 100644
--- a/ml/test.py
+++ b/ml/test.py
@@ -1,6 +1,8 @@
 import unittest
 from unittest.mock import Mock, patch
-from element_detection import detect, show_image_with_boxes, BoundBox
+from element_detection import detect, BoundBox
+from BlurMe.graphics.image_modification import show_image_with_boxes
+
 
 class TestYourModule(unittest.TestCase):