s352038/BlurMe
3
1
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

BLUR-37: Add detected object type to BoundingBox.

Browse Source
This commit is contained in:
Gosia 2023-12-19 09:28:06 +01:00
parent 07b7d24d40
commit 96da101cb4
1 changed files with 6 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
ml/element_detection.py
View File

@ -16,9 +16,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,14 +44,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