--- a
+++ b/YOLOv8.py
@@ -0,0 +1,60 @@
+"""
+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)