|
31 | 31 | from pathlib import Path |
32 | 32 | from typing import Any, Dict, Optional, cast |
33 | 33 |
|
| 34 | +try: |
| 35 | + from sharp.models import RGBGaussianPredictor |
| 36 | + |
| 37 | + ModelType = Optional[RGBGaussianPredictor] |
| 38 | +except ImportError: |
| 39 | + ModelType = Any |
| 40 | + |
34 | 41 | # 添加项目路径 |
35 | 42 | project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
36 | 43 | if project_root not in sys.path: |
@@ -70,7 +77,7 @@ def __init__(self, checkpoint_path: Optional[str] = None, device: str = "auto"): |
70 | 77 | "SHARP_MODEL_PATH", "models/sharp_2572gikvuh.pt" |
71 | 78 | ) |
72 | 79 | self.device = self._get_device(device) |
73 | | - self.model = None |
| 80 | + self.model: ModelType = None |
74 | 81 | self.loaded = False |
75 | 82 | self.load_time = 0.0 |
76 | 83 | self.metadata_file: Optional[str] = None |
@@ -142,17 +149,14 @@ def preload_model(self) -> bool: |
142 | 149 |
|
143 | 150 | # 创建和初始化模型 |
144 | 151 | self.model = create_predictor(PredictorParams()) |
145 | | - if self.model is not None: |
146 | | - self.model.load_state_dict(state_dict) |
147 | | - self.model.eval() |
148 | | - self.model.to(self.device) |
149 | | - |
150 | | - self.logger.info(f" 模型类型: {type(self.model).__name__}") |
151 | | - self.logger.info( |
152 | | - f" 模型参数量: {sum(p.numel() for p in self.model.parameters()):,}" |
153 | | - ) |
154 | | - else: |
155 | | - self.logger.warning("模型创建失败") |
| 152 | + self.model.load_state_dict(state_dict) |
| 153 | + self.model.eval() |
| 154 | + self.model.to(self.device) |
| 155 | + |
| 156 | + self.logger.info(f" 模型类型: {type(self.model).__name__}") |
| 157 | + self.logger.info( |
| 158 | + f" 模型参数量: {sum(p.numel() for p in self.model.parameters()):,}" |
| 159 | + ) |
156 | 160 |
|
157 | 161 | self.load_time = time.time() - start_time |
158 | 162 | self.loaded = True |
@@ -327,4 +331,3 @@ def main(): |
327 | 331 |
|
328 | 332 | if __name__ == "__main__": |
329 | 333 | main() |
330 | | - |
|
0 commit comments