Commit f3e7cfa7 authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): update the face

parent fa809d69

zoo/detection/face.py

0 → 100644
+26 −0
Original line number Diff line number Diff line 
import os.path



from ultralytics import YOLO



from ..base import _TRAIN_DIR as _GLOBAL_TRAIN_DIR



_TRAIN_DIR = os.path.join(_GLOBAL_TRAIN_DIR, 'face_detect')





def train(train_cfg: str, session_name: str, level: str = 's',

          max_epochs: int = 200, **kwargs):

    # Load a pretrained YOLO model (recommended for training)

    _last_pt = os.path.join(_TRAIN_DIR, session_name, 'weights', 'last.pt')

    if os.path.exists(_last_pt):

        model, resume = YOLO(_last_pt), True

    else:

        model, resume = YOLO(f'yolov8{level}.pt'), False



    # Train the model using the 'coco128.yaml' dataset for 3 epochs

    model.train(

        data=train_cfg, epochs=max_epochs,

        name=session_name, project=_TRAIN_DIR,

        save=True, plots=True,

        exist_ok=True, resume=resume,

        **kwargs

    )

zoo/detection/head.py

+2 −2
Original line number Diff line number Diff line
@@ -8,7 +8,7 @@ _TRAIN_DIR = os.path.join(_GLOBAL_TRAIN_DIR, 'head_detect') 




def train(train_cfg: str, session_name: str, level: str = 's',

          max_epochs: int = 300, save_per_epoch: int = 10, **kwargs):

          max_epochs: int = 200, **kwargs):

    # Load a pretrained YOLO model (recommended for training)

    _last_pt = os.path.join(_TRAIN_DIR, session_name, 'weights', 'last.pt')

    if os.path.exists(_last_pt):
@@ -20,7 +20,7 @@ def train(train_cfg: str, session_name: str, level: str = 's', 
    model.train(

        data=train_cfg, epochs=max_epochs,

        name=session_name, project=_TRAIN_DIR,

        save=True, save_period=save_per_epoch, plots=True,

        save=True, plots=True,

        exist_ok=True, resume=resume,

        **kwargs

    )

zoo/detection/manbits.py

+2 −2
Original line number Diff line number Diff line
@@ -8,7 +8,7 @@ _TRAIN_DIR = os.path.join(_GLOBAL_TRAIN_DIR, 'manbits_detect') 




def train(train_cfg: str, session_name: str, level: str = 'm',

          max_epochs: int = 300, save_per_epoch: int = 10, **kwargs):

          max_epochs: int = 200, **kwargs):

    # Load a pretrained YOLO model (recommended for training)

    _last_pt = os.path.join(_TRAIN_DIR, session_name, 'weights', 'last.pt')

    if os.path.exists(_last_pt):
@@ -20,7 +20,7 @@ def train(train_cfg: str, session_name: str, level: str = 'm', 
    model.train(

        data=train_cfg, epochs=max_epochs,

        name=session_name, project=_TRAIN_DIR,

        save=True, save_period=save_per_epoch, plots=True,

        save=True, plots=True,

        exist_ok=True, resume=resume,

        **kwargs

    )

zoo/detection/person.py

+2 −2
Original line number Diff line number Diff line
@@ -8,7 +8,7 @@ _TRAIN_DIR = os.path.join(_GLOBAL_TRAIN_DIR, 'person_detect') 




def train(train_cfg: str, session_name: str, level: str = 's',

          max_epochs: int = 300, save_per_epoch: int = 10, **kwargs):

          max_epochs: int = 200, **kwargs):

    # Load a pretrained YOLO model (recommended for training)

    _last_pt = os.path.join(_TRAIN_DIR, session_name, 'weights', 'last.pt')

    if os.path.exists(_last_pt):
@@ -20,7 +20,7 @@ def train(train_cfg: str, session_name: str, level: str = 's', 
    model.train(

        data=train_cfg, epochs=max_epochs,

        name=session_name, project=_TRAIN_DIR,

        save=True, save_period=save_per_epoch, plots=True,

        save=True, plots=True,

        exist_ok=True, resume=resume,

        **kwargs

    )