from app.core.segformer.detect import Detection as SegFormer, DetectionMock as SegFormerMock
from app.core.yolo.detect import YOLOSeg


class Model:
    def __init__(self):
        from app.main import MODEL
        if MODEL == "segformer":
            print("使用 SegFormer 模型")
            self.detection = SegFormer()
        elif MODEL == "yolo":
            print("使用 YOLO 模型")
            self.detection = YOLOSeg()

    def getModel(self):
        return self.detection