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
|
||||
|
||||
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)
|
||||
|
@ -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):
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user