fi0r1/roop
1
0
Fork 0
mirror of https://github.com/s0md3v/roop.git synced 2025-12-06 18:08:29 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Improve return typing

Browse Source
This commit is contained in:
henryruhs 2023-06-30 09:04:33 +02:00
parent eca4f0f127
commit b104741e68
2 changed files with 10 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
roop/capturer.py
View File

@ -1,8 +1,10 @@
from typing import Any
from typing import Optional
import cv2
from roop.typing import Frame
def get_video_frame(video_path: str, frame_number: int = 0) -> Any:
def get_video_frame(video_path: str, frame_number: int = 0) -> Optional[Frame]:
capture = cv2.VideoCapture(video_path)
frame_total = capture.get(cv2.CAP_PROP_FRAME_COUNT)
capture.set(cv2.CAP_PROP_POS_FRAMES, min(frame_total, frame_number - 1))

12
roop/face_analyser.py
View File

@ -1,9 +1,9 @@
import threading
from typing import Any
from typing import Any, Optional, List
import insightface
import roop.globals
from roop.typing import Frame
from roop.typing import Frame, Face
FACE_ANALYSER = None
THREAD_LOCK = threading.Lock()
@ -19,15 +19,15 @@ def get_face_analyser() -> Any:
return FACE_ANALYSER
def get_one_face(frame: Frame) -> Any:
face = get_face_analyser().get(frame)
def get_one_face(frame: Frame) -> Optional[Face]:
faces = get_many_faces(frame)
try:
return min(face, key=lambda x: x.bbox[0])
return faces[0]
except ValueError:
return None
def get_many_faces(frame: Frame) -> Any:
def get_many_faces(frame: Frame) -> Optional[List[Face]]:
try:
return get_face_analyser().get(frame)
except IndexError: