Compare commits
5 Commits
07b7d24d40
...
6d2b939987
Author | SHA1 | Date | |
---|---|---|---|
6d2b939987 | |||
b27f4930db | |||
e990582cbd | |||
9f55a0c034 | |||
96da101cb4 |
48
graphics/image_modification.py
Normal file
48
graphics/image_modification.py
Normal file
@ -0,0 +1,48 @@
|
|||||||
|
import os
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
from PIL import Image, ImageDraw, ImageFont, ImageFilter
|
||||||
|
from PIL.Image import composite
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
|
||||||
|
def blur_boxes_on_image(
|
||||||
|
in_image_path: str, bounding_boxes: List[BoundBox], out_image_path: str = None
|
||||||
|
):
|
||||||
|
img = Image.open(in_image_path)
|
||||||
|
blur_img = img.filter(ImageFilter.GaussianBlur(radius=7))
|
||||||
|
mask = Image.new("L", img.size, 255)
|
||||||
|
draw_mask = ImageDraw.Draw(mask)
|
||||||
|
for box in bounding_boxes:
|
||||||
|
if box.selected:
|
||||||
|
if box.object == "plate":
|
||||||
|
draw_mask.rectangle(box.get_params(), fill=0)
|
||||||
|
elif box.object == "face":
|
||||||
|
draw_mask.ellipse(box.get_params(), fill=0)
|
||||||
|
mask = mask.filter(ImageFilter.GaussianBlur(radius=3))
|
||||||
|
img = composite(img, blur_img, mask)
|
||||||
|
if not out_image_path:
|
||||||
|
out_image_path = (
|
||||||
|
in_image_path.split(".")[0] + "_blurred." + 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
|
||||||
|
|
||||||
@ -16,9 +15,12 @@ IOU_THRESH = 0.5
|
|||||||
|
|
||||||
|
|
||||||
class BoundBox:
|
class BoundBox:
|
||||||
def __init__(self, x1, y1, x2, y2):
|
def __init__(self, x1, y1, x2, y2, object=None):
|
||||||
self.x1, self.y1, self.x2, self.y2 = x1, y1, x2, y2
|
self.x1, self.y1, self.x2, self.y2 = x1, y1, x2, y2
|
||||||
self.selected = True
|
self.selected = True
|
||||||
|
if object not in ["face", "plate"]:
|
||||||
|
raise ValueError("object must be either 'face' or 'plate'")
|
||||||
|
self.object = object
|
||||||
|
|
||||||
def select(self):
|
def select(self):
|
||||||
self.selected = True
|
self.selected = True
|
||||||
@ -41,29 +43,14 @@ def detect(image_path: str) -> List[BoundBox]:
|
|||||||
)
|
)
|
||||||
plates = plates[0].cpu().numpy().boxes
|
plates = plates[0].cpu().numpy().boxes
|
||||||
bounding_boxes = []
|
bounding_boxes = []
|
||||||
for boxes in [faces, plates]:
|
for boxes, tag in zip([faces, plates], ["face", "plate"]):
|
||||||
for box in boxes:
|
for box in boxes:
|
||||||
xyxyn = box.xyxy[0]
|
xyxyn = box.xyxy[0]
|
||||||
x1 = int(xyxyn[0])
|
x1 = int(xyxyn[0])
|
||||||
y1 = int(xyxyn[1])
|
y1 = int(xyxyn[1])
|
||||||
x2 = int(xyxyn[2])
|
x2 = int(xyxyn[2])
|
||||||
y2 = int(xyxyn[3])
|
y2 = int(xyxyn[3])
|
||||||
bounding_boxes.append(BoundBox(x1, y1, x2, y2))
|
bounding_boxes.append(BoundBox(x1, y1, x2, y2, tag))
|
||||||
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