File size: 2,051 Bytes
17c5137 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
import os
import sys
import shutil
from ultralytics import YOLO
from src.logger import logging
from src.exception import PlantException
from src.entity.config_entity import ModelTrainerConfig
from src.entity.artifact_entity import ModelTrainerArtifact
from src.entity.artifact_entity import DataIngestionArtifact
class ModelTrainer:
try:
def __init__(self, model_trainer_config:ModelTrainerConfig, data_ingestion_artifact:DataIngestionArtifact):
self.model_trainer_config = model_trainer_config
self.data_ingestion_artifact = data_ingestion_artifact
def initiate_model_trainer(self) -> ModelTrainerArtifact:
logging.info(f"Removing and existing runs directory from previous training")
os.system("rm -rf runs")
model_config_file_name = self.model_trainer_config.weight_name.split('.')[0]
print(model_config_file_name)
os.system(f"yolo task=classify mode=train model={self.model_trainer_config.weight_name} \
data={self.data_ingestion_artifact.dataset_path} epochs={self.model_trainer_config.no_epochs} \
imgsz=128 batch={self.model_trainer_config.batch_size} patience={self.model_trainer_config.patience}")
os.makedirs("custom_model_weights", exist_ok=True)
os.system("cp runs/classify/train/weights/best.pt custom_model_weights/")
os.makedirs(self.model_trainer_config.model_trainer_dir, exist_ok=True)
os.system(f"cp runs/classify/train/weights/best.pt {self.model_trainer_config.model_trainer_dir}/")
model_trainer_artifact = ModelTrainerArtifact(trained_model_file_path="custom_model_weights/best.pt",)
logging.info("Exited initiate_model_trainer method of ModelTrainer class")
logging.info(f"Model trainer artifact: {model_trainer_artifact}")
return model_trainer_artifact
except Exception as e:
raise PlantException(e, sys) |