File size: 3,509 Bytes
3c87883 |
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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
# 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)
|