"""

data:

    CT         :     used

    mask       : not used

    labels(txt):     used

    labelsJson : not used

    dataset.yaml:    used

"""

from ultralytics import YOLO

import os

import shutil

import torch

import multiprocessing

if __name__ == "__main__":

    if torch.cuda.is_available():

        device = torch.device("cuda:0")

        print("GPU")

    else:

        device = torch.device("cpu")

        print("CPU")

    multiprocessing.freeze_support()

    torch.cuda.empty_cache()

    YOLODataset_path = "./YOLODataset"

    models_path = "./models/"

    save_result_path = "./runs"

    if os.path.exists(save_result_path):

        shutil.rmtree(save_result_path)

    # Train the model

    # Load a model

    yolov8_premodel_path = os.path.join(models_path, "yolov8n-seg.pt")

    yaml_dataset_path = os.path.join(YOLODataset_path, "dataset.yaml")

    model = YOLO(yolov8_premodel_path)  # load a pretrained model

    results = model.train(

        data=yaml_dataset_path,

        epochs=1,

        # imgsz=640,

        batch=2,

        flipud=0.5,

        # device="cpu",

        # device=0,

        seed=2024,

        amp=True,

    )

    model.val()

    model = YOLO("./runs/segment/train/weights/best.pt")

    source = "./test/testpng/"

    for f in os.listdir(source):

        file_path = os.path.join(source, f)

        model.predict(file_path, save=True, retina_masks=True)