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


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()
        elif MODEL == "yolo_detect":
            print("使用 YOLO_Detect 模型")
            self.detection = YOLODetect()

    def getModel(self):
        return self.detection