Update model_loaders.py
Browse files- model_loaders.py +25 -0
model_loaders.py
CHANGED
|
@@ -10,6 +10,7 @@
|
|
| 10 |
import streamlit as st
|
| 11 |
import torch
|
| 12 |
import psutil
|
|
|
|
| 13 |
|
| 14 |
logger = logging.getLogger(__name__)
|
| 15 |
|
|
@@ -96,6 +97,11 @@ def load_sam2_predictor():
|
|
| 96 |
logger.error(f"Failed to load SAM2 predictor: {e}")
|
| 97 |
return None
|
| 98 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 99 |
# ============================================================================
|
| 100 |
# MatAnyone Model Loading
|
| 101 |
# ============================================================================
|
|
@@ -119,6 +125,25 @@ def load_matanyone_processor():
|
|
| 119 |
logger.error(f"Failed to load MatAnyone: {e}")
|
| 120 |
return None
|
| 121 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 122 |
# ============================================================================
|
| 123 |
# Model Health Check
|
| 124 |
# ============================================================================
|
|
|
|
| 10 |
import streamlit as st
|
| 11 |
import torch
|
| 12 |
import psutil
|
| 13 |
+
import mediapipe as mp
|
| 14 |
|
| 15 |
logger = logging.getLogger(__name__)
|
| 16 |
|
|
|
|
| 97 |
logger.error(f"Failed to load SAM2 predictor: {e}")
|
| 98 |
return None
|
| 99 |
|
| 100 |
+
# Alias for new app.py
|
| 101 |
+
def load_sam2():
|
| 102 |
+
"""Alias for load_sam2_predictor() - for compatibility with app.py"""
|
| 103 |
+
return load_sam2_predictor()
|
| 104 |
+
|
| 105 |
# ============================================================================
|
| 106 |
# MatAnyone Model Loading
|
| 107 |
# ============================================================================
|
|
|
|
| 125 |
logger.error(f"Failed to load MatAnyone: {e}")
|
| 126 |
return None
|
| 127 |
|
| 128 |
+
# Alias for new app.py
|
| 129 |
+
def load_matanyone():
|
| 130 |
+
"""Alias for load_matanyone_processor() - for compatibility with app.py"""
|
| 131 |
+
return load_matanyone_processor()
|
| 132 |
+
|
| 133 |
+
# ============================================================================
|
| 134 |
+
# MediaPipe Pose
|
| 135 |
+
# ============================================================================
|
| 136 |
+
|
| 137 |
+
# Initialize MediaPipe Pose as a module-level variable
|
| 138 |
+
mp_pose = mp.solutions.pose
|
| 139 |
+
pose = mp_pose.Pose(
|
| 140 |
+
static_image_mode=False,
|
| 141 |
+
model_complexity=1,
|
| 142 |
+
enable_segmentation=True,
|
| 143 |
+
min_detection_confidence=0.5
|
| 144 |
+
)
|
| 145 |
+
logger.info("✅ MediaPipe Pose initialized")
|
| 146 |
+
|
| 147 |
# ============================================================================
|
| 148 |
# Model Health Check
|
| 149 |
# ============================================================================
|