Merge commit '9e8db857f4661c771d3e492a7641c2d189775daf' as 'utils_v2'

2024-12-24 11:22:28 +05:30
parent a13d66d106 9e8db857f4
commit bedf7da07d
168 changed files with 136836 additions and 0 deletions
@@ -0,0 +1,176 @@
+"""
+
+    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))
+