Commit 896d889f authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): add zoo.detection

parent f2ba719d

.gitignore

+4 −1
Original line number Diff line number Diff line
@@ -1218,3 +1218,6 @@ fabric.properties 
/test/testfile/dataset

/zoo/testfile/dataset

/.train

*.pt

/runs

/YOLOv8

requirements-zoo.txt

+2 −1
Original line number Diff line number Diff line
@@ -16,3 +16,4 @@ timm 
ftfy

regex

git+https://github.com/openai/CLIP.git

ultralytics

zoo/detection/__init__.py

0 → 100644
+0 −0

Empty file added.

zoo/detection/face.py

0 → 100644
+24 −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 = 300, **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,

        exist_ok=True, save=True, resume=resume,

        **kwargs

    )

zoo/detection/person.py

0 → 100644
+24 −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, 'person_detect')





def train(train_cfg: str, session_name: str, level: str = 's', max_epochs: int = 300, **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,

        exist_ok=True, save=True, resume=resume,

        **kwargs

    )