Spaces:
Sleeping
Sleeping
File size: 2,803 Bytes
36cd99b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
"""
Created By: ishwor subedi
Date: 2024-07-10
"""
import cv2
import numpy as np
from ultralytics import YOLO
class ObjectDetection:
def __init__(self, device="cpu"):
self.device = device
self.hand_detector = YOLO(model="artifacts/detection/hand_yolov8s.pt")
self.face_detector = YOLO(model="artifacts/detection/face_yolov8s.pt")
def hand_detection(self, image_path: str):
hand_detection_results = self.hand_detector.predict(image_path, show=False, device=self.device, iou=0.8)
orig_image = hand_detection_results[0].orig_img
boxes = hand_detection_results[0].boxes.xyxy #
return boxes, orig_image
def face_detection(self, image_path: str):
face_detection_results = self.face_detector.predict(image_path, show=False, device=self.device, iou=0.8)
boxes = face_detection_results[0].boxes.xyxy #
orig_image = face_detection_results[0].orig_img
return boxes, orig_image
def detect(self, hand: bool, face: bool, image_path: str):
if hand == True and face == True:
hand_boxes, orig_image = self.hand_detection(image_path)
face_boxes, orig_image = self.face_detection(image_path)
mask = np.zeros_like(orig_image)
for i in range(len(hand_boxes)):
box = hand_boxes[i].tolist()
x, y, x1, y1 = map(int, box[:4])
cv2.rectangle(mask, (x, y), (x1, y1), (255, 255, 255), -1)
print("hand")
for j in range(len(face_boxes)):
box = face_boxes[j].tolist()
x, y, x1, y1 = map(int, box[:4])
cv2.rectangle(mask, (x, y), (x1, y1), (255, 255, 255), -1)
print("face")
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
return mask
elif hand == True and face == False:
hand_boxes, orig_image = self.hand_detection(image_path)
mask = np.zeros_like(orig_image)
for i in range(len(hand_boxes)):
box = hand_boxes[i].tolist()
x, y, x1, y1 = map(int, box[:4])
cv2.rectangle(mask, (x, y), (x1, y1), (255, 255, 255), -1)
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
return mask
elif hand == False and face == True:
face_boxes, orig_image = self.face_detection(image_path)
mask = np.zeros_like(orig_image)
for j in range(len(face_boxes)):
box = face_boxes[j].tolist()
x, y, x1, y1 = map(int, box[:4])
cv2.rectangle(mask, (x, y), (x1, y1), (255, 255, 255), -1)
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
return mask
else:
return "none"
|