Merge branch 'master' of https://git.wmi.amu.edu.pl/s352038/BlurMe

graphics/image_modification.py

@@ -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)

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
 
@@ -16,9 +15,12 @@ IOU_THRESH = 0.5
 
 
 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.selected = True
+        if object not in ["face", "plate"]:
+            raise ValueError("object must be either 'face' or 'plate'")
+        self.object = object
 
     def select(self):
         self.selected = True
@@ -41,29 +43,14 @@ def detect(image_path: str) -> List[BoundBox]:
     )
     plates = plates[0].cpu().numpy().boxes
     bounding_boxes = []
-    for boxes in [faces, plates]:
+    for boxes, tag in zip([faces, plates], ["face", "plate"]):
         for box in boxes:
             xyxyn = box.xyxy[0]
             x1 = int(xyxyn[0])
             y1 = int(xyxyn[1])
             x2 = int(xyxyn[2])
             y2 = int(xyxyn[3])
-            bounding_boxes.append(BoundBox(x1, y1, x2, y2))
+            bounding_boxes.append(BoundBox(x1, y1, x2, y2, tag))
     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)

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):
 

ui.lnf.xml

@@ -0,0 +1,5 @@
+<application>
+  <component name="UISettings">
+    <option name="HIDE_TOOL_STRIPES" value="false" />
+  </component>
+</application>