khushalps
eef89c9ebe
git-subtree-dir: utils_v2 git-subtree-split: a6614afbe332e89b495f068705f04f085f931adf
177 lines
6.4 KiB
Python
177 lines
6.4 KiB
Python
"""
|
|
|
|
AUTHOR:
|
|
|
|
Khushal P Soonderji
|
|
|
|
DATE:
|
|
|
|
Thursday, 4th Jul, 2024
|
|
|
|
OBJECTIVE:
|
|
|
|
To provide a class to detect objects in images using YOLO behaviour_models.
|
|
|
|
REFERENCES:
|
|
|
|
1) Code Examples: https://docs.ultralytics.com/usage/python/
|
|
2) Models: https://docs.ultralytics.com/models/yolov8/#supported-tasks-and-modes
|
|
|
|
DOWNLOADS:
|
|
|
|
N/A
|
|
|
|
"""
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** IMPORT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# To make sibling directories accessible for imports:
|
|
import sys
|
|
sys.path.append(".")
|
|
sys.path.append("..")
|
|
|
|
# System-level activities:
|
|
import os
|
|
|
|
# To use YOLO architecture:
|
|
from ultralytics import YOLO
|
|
|
|
# To work with images:
|
|
from PIL import Image
|
|
|
|
# utils:
|
|
from utils_v2.string import json
|
|
|
|
# Common:
|
|
from shared import constants
|
|
|
|
# For debugging:
|
|
from icecream import IceCreamDebugger
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MACROS / ONE-TIME INIT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** VARIABLES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** FUNCTIONS ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
class YoloDetect:
|
|
|
|
def __init__(
|
|
self,
|
|
model_file,
|
|
debug = True,
|
|
debug_prefix = "YOLO | "
|
|
):
|
|
|
|
# Load the model:
|
|
self._model = YOLO(model = model_file)
|
|
|
|
# Initialize the dbugging tool:
|
|
self._printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True)
|
|
if not debug: self._printer.disable()
|
|
|
|
def disable_debug(self):
|
|
self._printer.disable()
|
|
|
|
def enable_debug(self):
|
|
self._printer.enable()
|
|
|
|
@staticmethod
|
|
def _process_result(raw_result):
|
|
|
|
# Extract the various types of results available from the inference:
|
|
speed = {k: (v / 100.0) for k, v in raw_result.speed.items()}
|
|
class_mapping = raw_result.names
|
|
probabilities = raw_result.probs
|
|
boxes = raw_result.boxes
|
|
masks = raw_result.masks
|
|
|
|
# Construct a default response:
|
|
response_json = {
|
|
"speed": speed,
|
|
"classes": class_mapping,
|
|
"boxes": None
|
|
}
|
|
|
|
# Object-Detection results:
|
|
formatted_boxes = []
|
|
for box in boxes:
|
|
detected_class = int(box.cls.numpy()[0])
|
|
x1, y1, x2, y2 = box.xyxy[0]
|
|
formatted_boxes.append(
|
|
{
|
|
"class": detected_class,
|
|
"className": class_mapping[detected_class],
|
|
"confidence": float(box.conf.numpy()[0]),
|
|
"x1": int(x1.numpy()),
|
|
"y1": int(y1.numpy()),
|
|
"x2": int(x2.numpy()),
|
|
"y2": int(y2.numpy())
|
|
}
|
|
)
|
|
response_json["boxes"] = formatted_boxes
|
|
|
|
# Done here:
|
|
return response_json
|
|
|
|
def predict(self, image, show = False, verbose = False):
|
|
|
|
# Process the input image with the given task:
|
|
results = self._model(
|
|
source = image,
|
|
show = show,
|
|
verbose = verbose
|
|
)
|
|
|
|
# Based on the task, interpret the results:
|
|
results_json = self._process_result(results[0])
|
|
|
|
# Done here:
|
|
return results_json
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MAIN PROGRAM ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
# model_file_path = os.path.join(constants.PROJECT_DIRECTORY, "ai", "yolo", "behaviour_models", "yolov8x.pt")
|
|
model_file_path = r"/home/developer/PycharmProjects/utils/data/ai/models/hugging_face/object_detection/YOLOv10-Document-Layout-Analysis/yolov10x_best.pt"
|
|
# sample_image_path = r"/home/developer/Downloads/2_cats.jpg"
|
|
sample_image_path = r"/home/developer/Downloads/flattened_image.jpg"
|
|
my_yolo = YoloDetect(model_file = model_file_path)
|
|
results = my_yolo.predict(image = Image.open(sample_image_path))
|
|
print("FINAL RESULTS:", json.to_string(results))
|
|
|