activity_detection/activity_detection.py

# this file containes class for Activity Detection.
from ultralytics import YOLO
import logging

# Configure the logger
logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s - %(levelname)s - %(message)s",
    filename="logs.log",
)

# Create a logger
logger = logging.getLogger("pipline")

class ActivityDetection:
    """Class to detect activity of person from image"""

    def __init__(self):
        self.activity_classes = ['Standing','Running','Sitting']
        self.input_file_path = None
        self.trained_model_path = None
        self.trained_model = None
        self.base_model = None
    
    def set_input_file_path(self, input_file_path):
        """
        Method to set path of input files to train a model

        Args:
            input_file_path (str): Relative or global path to the input files
        """
        self.input_file_path = input_file_path
        logger.info("input file path is set...")
        
    def set_trained_model_path(self, trained_model_path):
        """ 
        Method to set path of trained model to inference the model

        Args:
            trained_model_path (str): Relative or global path to the trained model
        """
        self.trained_model_path = trained_model_path
        logger.info("trained_model_path is set...")
        self.trained_model = YOLO(trained_model_path)
        logger.info("tranined_model is created successfully ...")

    def train(self):
        """
        Method to train a model for activity detection.

        Raises:
            BaseException: It generates BaseException when it could'f find input_file_path
            e: any other exception that occors.
        """
        self.base_model = YOLO("yolov8n-cls.pt")
        try:
            if self.input_file_path is None:
                raise BaseException("Please set path of input_files first with set_input_file_path method.")
            
            self.base_model.train(data=self.input_file_path, epochs=50)
            logger.info("traning of model is successfully done...")
        except Exception as e:
            logger.error("Something went wrong in activity detection model training")
            logger.error(e)
        
    def inference(self, single_object_images):
        """Method to detect activity of person from image.

        Args:
            single_object_images (list of numpy array): list of single object images

        Raises:
            BaseException: It generates BaseException when it could'f find trained_model_path
            e: any other exception that occors.

        Returns:
            activities(list of strings): returns list of activities perform by person
        """
        logger.info("inference method is called...")
        try:
            if self.trained_model is None:
                raise BaseException("Please set path of trained model first with set_trained_model_path method.")
            
            activities = []
            
            # detect activity in image
            for img in single_object_images:
                predictions =self.trained_model.predict(img)

                for result in predictions:

                    probs = result.probs
                    class_index = probs.top1

                activities.append(self.activity_classes[class_index])

            return activities
        
        except Exception as e:
            logger.error("Something went wrong in activity detection model inference")
            logger.error(e)