config.py

from cv2 import INTER_CUBIC

import plasma.functional as f


class Config(f.Pipe):

    def __init__(self) -> None:
        super().__init__(
            DEVICE = 'cuda:1',
            
            STANDARD_WIDTH = 1024,
            STANDARD_INTERPOLATION = INTER_CUBIC,
            TOGRAY = True,

            TABLE_DETECTION_CHECKPOINT = 'MedCat/checkpoints:table_detection_yolo_v3.pt',
            CONF_THRESHOLD = 0.8,
            YOLO_IMAGE_SIZE = 640,
            HEIGHT_EXPAND_RATIO = 0.03,

            USE_STAMP_DETECTION = True,
            STAMP_DETECTION_CHECKPOINT = 'MedCat/checkpoints:stamp_detection_yolo_v1.pt',
            ORIGINAL_STAMP_WEIGHT_PATH = 'https://drive.google.com/file/d/1YHH7pLoZEdyxw2AoLz9G4lrq6uuxweYB/view?usp=sharing',
            STAMP_REMOVING_IOU_THRESHOLD = 0.6
        )