Spaces:
Sleeping
Sleeping
| """ | |
| 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" | |