File size: 2,773 Bytes
22b4354 2b7fd54 71b7d31 6058d6e a83a3d9 e492152 04b83e7 f2aa528 ce9c587 cb528cc f2a34e7 cb528cc 6f578f7 ccfe1b7 6f578f7 f2aa528 6736291 22b4354 a83a3d9 22b4354 71b7d31 2b7fd54 f2aa528 a83a3d9 e492152 6058d6e 22b4354 |
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 |
import functools
import os
PROJECT_ROOT_DIR = os.path.join(os.path.abspath(os.path.dirname(__file__)), "..", "..")
MODELS_DIR = os.path.join(PROJECT_ROOT_DIR, "models")
MODELS_ANIMAL_DIR = os.path.join(MODELS_DIR, "animal")
OUTPUTS_DIR = os.path.join(PROJECT_ROOT_DIR, "outputs")
OUTPUTS_VIDEOS_DIR = os.path.join(OUTPUTS_DIR, "videos")
TEMP_DIR = os.path.join(OUTPUTS_DIR, "temp")
TEMP_VIDEO_FRAMES_DIR = os.path.join(TEMP_DIR, "video_frames")
TEMP_VIDEO_OUT_FRAMES_DIR = os.path.join(TEMP_VIDEO_FRAMES_DIR, "out")
EXP_OUTPUT_DIR = os.path.join(OUTPUTS_DIR, "exp_data")
MODEL_CONFIG = os.path.join(PROJECT_ROOT_DIR, "modules", "config", "models.yaml")
MODEL_PATHS = {
"appearance_feature_extractor": os.path.join(MODELS_DIR, "appearance_feature_extractor.safetensors"),
"motion_extractor": os.path.join(MODELS_DIR, "motion_extractor.safetensors"),
"warping_module": os.path.join(MODELS_DIR, "warping_module.safetensors"),
"spade_generator": os.path.join(MODELS_DIR, "spade_generator.safetensors"),
"stitching_retargeting_module": os.path.join(MODELS_DIR, "stitching_retargeting_module.safetensors"),
"face_yolov8n": os.path.join(MODELS_DIR, "face_yolov8n.pt")
}
MODEL_ANIMAL_PATHS = {
"appearance_feature_extractor": os.path.join(MODELS_ANIMAL_DIR, "appearance_feature_extractor.safetensors"),
"motion_extractor": os.path.join(MODELS_ANIMAL_DIR, "motion_extractor.safetensors"),
"warping_module": os.path.join(MODELS_ANIMAL_DIR, "warping_module.safetensors"),
"spade_generator": os.path.join(MODELS_ANIMAL_DIR, "spade_generator.safetensors"),
"stitching_retargeting_module": os.path.join(MODELS_ANIMAL_DIR, "stitching_retargeting_module.safetensors"),
# Just animal detection model not the face, needs better model
"yolo_v5s_animal_det": os.path.join(MODELS_ANIMAL_DIR, "yolo_v5s_animal_det.n2x")
}
MASK_TEMPLATES = os.path.join(PROJECT_ROOT_DIR, "modules", "utils", "resources", "mask_template.png")
I18N_YAML_PATH = os.path.join(PROJECT_ROOT_DIR, "i18n", "translation.yaml")
def get_auto_incremental_file_path(dir_path: str, extension: str, prefix: str = ""):
counter = 0
while True:
if prefix:
filename = f"{prefix}_{counter:05d}.{extension}"
else:
filename = f"{counter:05d}.{extension}"
full_path = os.path.join(dir_path, filename)
if not os.path.exists(full_path):
return full_path
counter += 1
@functools.lru_cache
def init_dirs():
for dir_path in [
MODELS_DIR,
MODELS_ANIMAL_DIR,
OUTPUTS_DIR,
EXP_OUTPUT_DIR,
TEMP_DIR,
TEMP_VIDEO_FRAMES_DIR,
TEMP_VIDEO_OUT_FRAMES_DIR,
OUTPUTS_VIDEOS_DIR
]:
os.makedirs(dir_path, exist_ok=True)
init_dirs()
|