models_loader.py 618 B

1234567891011121314151617181920
  1. from app.utils.MyBatchOnnxYolo import MyBatchOnnxYolo
  2. from app.utils.MyViTFeatureExtractor import MyViTFeatureExtractor
  3. from app.core.config import settings
  4. class GlobalModels:
  5. yolo: MyBatchOnnxYolo = None
  6. vit: MyViTFeatureExtractor = None
  7. @classmethod
  8. def load_models(cls):
  9. print("⏳ Loading YOLO Model...")
  10. cls.yolo = MyBatchOnnxYolo(settings.YOLO_MODEL_PATH, task='segment') # 或 detect
  11. print("⏳ Loading ViT Model...")
  12. cls.vit = MyViTFeatureExtractor(settings.VIT_MODEL_PATH)
  13. print("✅ Models Loaded Successfully.")
  14. global_models = GlobalModels()