pokouqiege/models/line_detect/train_demo.py

import torch
import os

from models.line_detect.line_detect import linedetect_newresnet18fpn, linedetect_resnet50_fpn, linedetect_resnet18_fpn, \
    linedetect_newresnet50fpn, linedetect_maxvitfpn, linedetect_high_maxvitfpn, linedetect_swin_transformer_fpn, \
    linedetect_newresnet101fpn, linedetect_newresnet152fpn, linedetect_efficientnet

from models.line_net.trainer import Trainer

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if __name__ == '__main__':

    os.environ["CUDA_LAUNCH_BLOCKING"] = "1"

    # model = LineNet('line_net.yaml')
    # model=linedetect_resnet50_fpn()
    # model = linedetect_resnet50_fpn()
    # model=get_line_net_convnext_fpn(num_classes=2).to(device)
    # model=linenet_newresnet50fpn()
    # model = lineDetect_resnet18_fpn()

    # model=linedetect_resnet18_fpn()
    # model=linedetect_newresnet18fpn(num_points=4)
    # model=linedetect_newresnet50fpn(num_points=4)
    # model = linedetect_newresnet101fpn(num_points=4)
    # model = linedetect_newresnet152fpn(num_points=4)
    # model.load_weights(save_path=r'/home/admin/projects/MultiVisionModels/models/line_detect/train_results/20250711_114046/weights/best_val.pth')
    # model=linedetect_maxvitfpn()

    model=linedetect_efficientnet(name='efficientnet_v2_l')
    # model=linedetect_high_maxvitfpn()

    # model=linedetect_swin_transformer_fpn(type='t')
    model.start_train(cfg='train.yaml')