BLUR-37: Refactor code to move creation of image with boxes to graphics directory.
This commit is contained in:
parent
96da101cb4
commit
9f55a0c034
26
graphics/add_boxes_to_image.py
Normal file
26
graphics/add_boxes_to_image.py
Normal file
@ -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)
|
@ -2,7 +2,6 @@ import os
|
|||||||
from typing import List, Tuple
|
from typing import List, Tuple
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
from PIL import Image, ImageDraw, ImageFont
|
|
||||||
|
|
||||||
from ultralytics import YOLO
|
from ultralytics import YOLO
|
||||||
|
|
||||||
@ -55,18 +54,3 @@ def detect(image_path: str) -> List[BoundBox]:
|
|||||||
return bounding_boxes
|
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)
|
|
||||||
|
@ -1,6 +1,8 @@
|
|||||||
import unittest
|
import unittest
|
||||||
from unittest.mock import Mock, patch
|
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):
|
class TestYourModule(unittest.TestCase):
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user