tuoheng_alg_move_20240715/config/ModelConfig.py

# -*- coding: utf-8 -*-
import sys

sys.path.extend(['..', '../AIlib'])
from segutils.segmodel import SegModel, get_largest_contours
from utils.torch_utils import select_device
from models.experimental import attempt_load
from utilsK.queRiver import get_labelnames, get_label_arrays, post_process_
from AI import AI_process, AI_process_forest, get_postProcess_para


class ModelConfig():
    def __init__(self):
        postFile = '../AIlib/conf/para.json'
        self.conf_thres, self.iou_thres, self.classes, self.rainbows = get_postProcess_para(postFile)


class SZModelConfig(ModelConfig):

    def __init__(self):
        super(SZModelConfig, self).__init__()
        labelnames = "../AIlib/weights/yolov5/class8/labelnames.json" ##对应类别表
        self.names = get_labelnames(labelnames)
        self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=40,
                                                fontpath="../AIlib/conf/platech.ttf")


class LCModelConfig(ModelConfig):
    def __init__(self):
        super(LCModelConfig, self).__init__()
        labelnames = "../AIlib/weights/forest/labelnames.json"
        self.names = get_labelnames(labelnames)
        self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=40, fontpath="../AIlib/conf/platech.ttf")


class RFModelConfig(ModelConfig):
    def __init__(self):
        super(RFModelConfig, self).__init__()
        labelnames = "../AIlib/weights/road/labelnames.json"
        self.names = get_labelnames(labelnames)
        imageW = 1536
        outfontsize=int(imageW/1920*40)
        self.label_arraylist = get_label_arrays(self.names, self.rainbows, outfontsize=outfontsize, fontpath="../AIlib/conf/platech.ttf")