9011167809

Committed 09 May 2024 01:56AM CUT coverage: 65.763%. First build

Build # 9011167809

Build Type

Pull #25

github

Committed by

web-flow

Commit Message

Merge 7c160e431 into 23491494b

Pull Request Pull Request #25: hotfix: add lightning dependancy for _cosine_scheduler_fn(..)

Run Details

461 of 701 relevant lines covered (65.76%)

0.66 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

19.64

/src/marlin_pytorch/face_detector.py

import glob
import os.path
import sys
from typing import Optional, Union, Tuple

import cv2
from numpy import ndarray

from .util import read_yaml, Singleton, crop_with_padding


@Singleton
class FaceXZooFaceDetector:

    def __init__(self):
        self.faceDetModelHandler = None
        self.inited = False

    def init(self, face_sdk_path: Optional[str] = None, device: str = "cuda:0"):
        if face_sdk_path is not None:
            sys.path.append(face_sdk_path)
        else:
            if os.path.exists("FaceX-Zoo"):
                face_sdk_path = "FaceX-Zoo"
                sys.path.append(face_sdk_path)
        try:
            from core.model_handler.face_detection.FaceDetModelHandler import FaceDetModelHandler
            from core.model_loader.face_detection.FaceDetModelLoader import FaceDetModelLoader
        except ImportError:
            raise ImportError("FaceX-Zoo cannot be imported, please specify the path to the face_sdk path of FaceXZoo"
                              " or put it in the working directory.")

        model_conf = read_yaml(os.path.join(face_sdk_path, "config", "model_conf.yaml"))
        model_path = os.path.join(face_sdk_path, 'models')
        scene = 'non-mask'
        model_category = 'face_detection'
        model_name = model_conf[scene][model_category]

        faceDetModelLoader = FaceDetModelLoader(model_path, model_category, model_name)
        model, cfg = faceDetModelLoader.load_model()
        self.faceDetModelHandler = FaceDetModelHandler(model, device, cfg)
        self.inited = True

    @staticmethod
    def install(path: Optional[str] = None) -> str:
        """
        Install FaceX-Zoo by clone from GitHub.

        Args:
            path (``str``, optional): The path to install FaceX-Zoo, default is "./FaceX-Zoo".

        Returns:
            ``str``: The path to the installed FaceX-Zoo.

        """
        path = path or "FaceX-Zoo"
        if os.path.exists(path):
            return path

        os.system(f"git clone --depth=1 https://github.com/ControlNet/FaceX-Zoo {path or ''}")
        return path

    def detect_face(self, image: ndarray):
        assert image.ndim == 3 and image.shape[2] == 3, "frame should be 3-dim"
        dets = self.faceDetModelHandler.inference_on_image(image)
        return dets

    def crop_face(self, frame, margin=1, x=0, y=0) -> Tuple[ndarray, int, int, int]:
        dets = self.detect_face(frame)
        if len(dets) > 0:
            x1, y1, x2, y2, confidence = dets[0]
            # center
            x, y = (int((x1 + x2) / 2), int((y1 + y2) / 2))
            margin = int(max(abs(x2 - x1), abs(y2 - y1)) / 2)
        # crop face
        face = crop_with_padding(frame, x - margin, x + margin, y - margin, y + margin, 0)
        face = cv2.resize(face, (224, 224))
        return face, margin, x, y

    def crop_image(self, image_path: str, out_path: str, max_faces=1, margin=0) -> None:
        if max_faces > 1:
            raise NotImplementedError("Multiple faces are not supported yet.")

        frame = cv2.imread(image_path)
        dets = self.detect_face(frame)
        for det in dets[:max_faces]:
            x1, y1, x2, y2, _ = det
            cropped = crop_with_padding(frame, int(x1 - margin), int(x2 + margin), int(y1 - margin), int(y2 + margin))
            cv2.imwrite(out_path, cropped)

    def crop_video(self, video_path: str, out_path: str, frame_size: Optional[Union[int, Tuple[int, int]]] = None,
        margin=0, fourcc="mp4v"
    ) -> None:
        video = cv2.VideoCapture(video_path)
        if not video.isOpened():
            raise IOError("Cannot open video file: " + video_path)

        # infer frame size
        if frame_size is None:
            frame_size = self._infer_frame_size(video_path, margin)

        if type(frame_size) is int:
            frame_size = frame_size, frame_size

        fps = video.get(cv2.CAP_PROP_FPS)
        writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*fourcc), fps, frame_size)
        x1, y1, x2, y2 = 0, 0, 1, 1
        while True:
            ret, frame = video.read()
            if not ret:
                break
            dets = self.detect_face(frame)
            if len(dets) > 0:
                x1, y1, x2, y2, confidence = dets[0]
                # center
                x, y = int((x1 + x2) / 2), int((y1 + y2) / 2)
                side = int(max(abs(x2 - x1), abs(y2 - y1)))
                x1 = x - side // 2
                x2 = x + side // 2
                y1 = y - side // 2
                y2 = y + side // 2

            cropped = crop_with_padding(frame, int(x1 - margin), int(x2 + margin), int(y1 - margin), int(y2 + margin))
            resized = cv2.resize(cropped, frame_size)
            writer.write(resized)
        video.release()
        writer.release()

    def crop_image_dir(self, image_dir: str, out_dir: str, pattern="*.jpg", *args, **kwargs) -> None:
        all_images = glob.glob(os.path.join(image_dir, pattern), root_dir=image_dir)
        for image_path in all_images:
            out_path = os.path.join(out_dir, image_path)
            self.crop_image(image_path, out_path, *args, **kwargs)

    def crop_video_dir(self, video_dir: str, out_dir: str, pattern="*.mp4", *args, **kwargs) -> None:
        all_videos = glob.glob(os.path.join(video_dir, pattern), root_dir=video_dir)
        for video_path in all_videos:
            out_path = os.path.join(out_dir, video_path)
            self.crop_video(video_path, out_path, *args, **kwargs)

    def _infer_frame_size(self, video_path: str, margin: int = 0
    ) -> Tuple[int, int]:
        video = cv2.VideoCapture(video_path)
        while True:
            ret, frame = video.read()
            if not ret:
                break
            dets = self.detect_face(frame)
            if len(dets) > 0:
                x1, y1, x2, y2, confidence = dets[0]
                # center
                side = int(max(abs(x2 - x1), abs(y2 - y1)))
                video.release()
                return side + 2 * margin, side + 2 * margin

1	import glob	1✔
2	import os.path	1✔
3	import sys	1✔
4	from typing import Optional, Union, Tuple	1✔
5
6	import cv2	1✔
7	from numpy import ndarray	1✔
8
9	from .util import read_yaml, Singleton, crop_with_padding	1✔
10
11
12	@Singleton	1✔
13	class FaceXZooFaceDetector:	1✔
14
15	def __init__(self):	1✔
16	self.faceDetModelHandler = None	1✔
17	self.inited = False	1✔
18
19	def init(self, face_sdk_path: Optional[str] = None, device: str = "cuda:0"):	1✔
20	if face_sdk_path is not None:	×
21	sys.path.append(face_sdk_path)	×
22	else:
23	if os.path.exists("FaceX-Zoo"):	×
24	face_sdk_path = "FaceX-Zoo"	×
25	sys.path.append(face_sdk_path)	×
26	try:	×
27	from core.model_handler.face_detection.FaceDetModelHandler import FaceDetModelHandler	×
28	from core.model_loader.face_detection.FaceDetModelLoader import FaceDetModelLoader	×
29	except ImportError:	×
30	raise ImportError("FaceX-Zoo cannot be imported, please specify the path to the face_sdk path of FaceXZoo"	×
31	" or put it in the working directory.")
32
33	model_conf = read_yaml(os.path.join(face_sdk_path, "config", "model_conf.yaml"))	×
34	model_path = os.path.join(face_sdk_path, 'models')	×
35	scene = 'non-mask'	×
36	model_category = 'face_detection'	×
37	model_name = model_conf[scene][model_category]	×
38
39	faceDetModelLoader = FaceDetModelLoader(model_path, model_category, model_name)	×
40	model, cfg = faceDetModelLoader.load_model()	×
41	self.faceDetModelHandler = FaceDetModelHandler(model, device, cfg)	×
42	self.inited = True	×
43
44	@staticmethod	1✔
45	def install(path: Optional[str] = None) -> str:	1✔
46	"""
47	Install FaceX-Zoo by clone from GitHub.
48
49	Args:
50	path (``str``, optional): The path to install FaceX-Zoo, default is "./FaceX-Zoo".
51
52	Returns:
53	``str``: The path to the installed FaceX-Zoo.
54
55	"""
56	path = path or "FaceX-Zoo"	×
57	if os.path.exists(path):	×
58	return path	×
59
60	os.system(f"git clone --depth=1 https://github.com/ControlNet/FaceX-Zoo {path or ''}")	×
61	return path	×
62
63	def detect_face(self, image: ndarray):	1✔
64	assert image.ndim == 3 and image.shape[2] == 3, "frame should be 3-dim"	×
65	dets = self.faceDetModelHandler.inference_on_image(image)	×
66	return dets	×
67
68	def crop_face(self, frame, margin=1, x=0, y=0) -> Tuple[ndarray, int, int, int]:	1✔
69	dets = self.detect_face(frame)	×
70	if len(dets) > 0:	×
71	x1, y1, x2, y2, confidence = dets[0]	×
72	# center
73	x, y = (int((x1 + x2) / 2), int((y1 + y2) / 2))	×
74	margin = int(max(abs(x2 - x1), abs(y2 - y1)) / 2)	×
75	# crop face
76	face = crop_with_padding(frame, x - margin, x + margin, y - margin, y + margin, 0)	×
77	face = cv2.resize(face, (224, 224))	×
78	return face, margin, x, y	×
79
80	def crop_image(self, image_path: str, out_path: str, max_faces=1, margin=0) -> None:	1✔
81	if max_faces > 1:	×
82	raise NotImplementedError("Multiple faces are not supported yet.")	×
83
84	frame = cv2.imread(image_path)	×
85	dets = self.detect_face(frame)	×
86	for det in dets[:max_faces]:	×
87	x1, y1, x2, y2, _ = det	×
88	cropped = crop_with_padding(frame, int(x1 - margin), int(x2 + margin), int(y1 - margin), int(y2 + margin))	×
89	cv2.imwrite(out_path, cropped)	×
90
91	def crop_video(self, video_path: str, out_path: str, frame_size: Optional[Union[int, Tuple[int, int]]] = None,	1✔
92	margin=0, fourcc="mp4v"
93	) -> None:
94	video = cv2.VideoCapture(video_path)	×
95	if not video.isOpened():	×
96	raise IOError("Cannot open video file: " + video_path)	×
97
98	# infer frame size
99	if frame_size is None:	×
100	frame_size = self._infer_frame_size(video_path, margin)	×
101
102	if type(frame_size) is int:	×
103	frame_size = frame_size, frame_size	×
104
105	fps = video.get(cv2.CAP_PROP_FPS)	×
106	writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*fourcc), fps, frame_size)	×
107	x1, y1, x2, y2 = 0, 0, 1, 1	×
108	while True:	×
109	ret, frame = video.read()	×
110	if not ret:	×
111	break	×
112	dets = self.detect_face(frame)	×
113	if len(dets) > 0:	×
114	x1, y1, x2, y2, confidence = dets[0]	×
115	# center
116	x, y = int((x1 + x2) / 2), int((y1 + y2) / 2)	×
117	side = int(max(abs(x2 - x1), abs(y2 - y1)))	×
118	x1 = x - side // 2	×
119	x2 = x + side // 2	×
120	y1 = y - side // 2	×
121	y2 = y + side // 2	×
122
123	cropped = crop_with_padding(frame, int(x1 - margin), int(x2 + margin), int(y1 - margin), int(y2 + margin))	×
124	resized = cv2.resize(cropped, frame_size)	×
125	writer.write(resized)	×
126	video.release()	×
127	writer.release()	×
128
129	def crop_image_dir(self, image_dir: str, out_dir: str, pattern=".jpg", args, **kwargs) -> None:	1✔
130	all_images = glob.glob(os.path.join(image_dir, pattern), root_dir=image_dir)	×
131	for image_path in all_images:	×
132	out_path = os.path.join(out_dir, image_path)	×
133	self.crop_image(image_path, out_path, args, *kwargs)	×
134
135	def crop_video_dir(self, video_dir: str, out_dir: str, pattern=".mp4", args, **kwargs) -> None:	1✔
136	all_videos = glob.glob(os.path.join(video_dir, pattern), root_dir=video_dir)	×
137	for video_path in all_videos:	×
138	out_path = os.path.join(out_dir, video_path)	×
139	self.crop_video(video_path, out_path, args, *kwargs)	×
140
141	def _infer_frame_size(self, video_path: str, margin: int = 0	1✔
142	) -> Tuple[int, int]:
143	video = cv2.VideoCapture(video_path)	×
144	while True:	×
145	ret, frame = video.read()	×
146	if not ret:	×
147	break	×
148	dets = self.detect_face(frame)	×
149	if len(dets) > 0:	×
150	x1, y1, x2, y2, confidence = dets[0]	×
151	# center
152	side = int(max(abs(x2 - x1), abs(y2 - y1)))	×
153	video.release()	×
154	return side + 2 * margin, side + 2 * margin	×

ControlNet / MARLIN / 9011167809

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous