From 0ddb524ff544532a4b2d36ef6b4d3913ace0869a Mon Sep 17 00:00:00 2001 From: Daniel Parra Date: Sun, 24 Mar 2024 16:58:45 -0400 Subject: [PATCH 1/2] Began working on creating the Object detection layer --- layers/ObjectDetectionLayer.py | 138 +++++++++++++++++++++++++++++++++ 1 file changed, 138 insertions(+) create mode 100644 layers/ObjectDetectionLayer.py diff --git a/layers/ObjectDetectionLayer.py b/layers/ObjectDetectionLayer.py new file mode 100644 index 0000000..cc7039f --- /dev/null +++ b/layers/ObjectDetectionLayer.py @@ -0,0 +1,138 @@ +import cv2 + +import os + +import numpy as np +import torch +import torch.nn as nn +import tensorflow as tf +from PIL import Image +from torchvision import transforms +from ..ml.yolov7.models.experimental import attempt_load +from ..ml.yolov7.utils.plots import plot_one_box +from ..ml.yolov7.utils.general import non_max_suppression + +from Layer import AnalysisLayer + +class ObjectDetectionLayer(AnalysisLayer): + def __init__(self, in_size, out_size, path_to_weights, conf, classes_array, colors_array) -> None: + super().__init__(in_size, out_size, "objectDetection") + + class PTWeightsProcessor(): + def __init__(self, weights_path, conf_thres, iou_thres) -> None: + self.weights_path = weights_path + self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + + self.conf_thres = conf_thres + self.iou_thres = iou_thres + + absolute_file_path = os.path.abspath( + os.path.join(os.path.dirname(__file__), self.weights_path), + ) + + self.__MODEL = attempt_load( + absolute_file_path, + map_location=self.device, + ) + + def process(self, img): + image = Image.fromarray(img) + img_transform = transforms.Compose([transforms.ToTensor()]) + img_tensor = img_transform(image).to(self.device).unsqueeze(0) + + pred_results = self.__MODEL(img_tensor)[0] + detections = non_max_suppression( + pred_results, + conf_thres=self.conf_thres, + iou_thres=self.iou_thres, + ) + + arr_image = np.array(img) + + if detections: + detections = detections[0] + for x1, y1, x2, y2, conf, cls in detections: + class_index = int(cls.cpu().item()) + print(f"{self.__CLASSES[class_index]} => {conf}") + plot_one_box( + [x1, y1, x2, y2], + arr_image, + label=f"{self.__CLASSES[class_index]}", + color=self.__COLORS[class_index], + line_thickness=2, + ) + + return arr_image + + class TFLiteWeightsProcessor(): + def __init__(self, weights_path) -> None: + self.weights_path = weights_path + + def process(self, img): + COLORS = [(255,0,0),(0,255,0),(0,0,255),(255,155,0),(255,0,255),(0,255,255)] + CLASSES = ['buoy_abydos_serpenscaput', 'buoy_abydos_taurus', 'buoy_earth_auriga', 'buoy_earth_cetus', 'gate_abydos', 'gate_earth'] + CONFIDENCE_THRESHOLD = 0.77 + + # Load the TFLite model and allocate tensors + interpreter = tf.lite.Interpreter(model_path=self.weights_path) + interpreter.allocate_tensors() + + + # Get input and output tensors + input_details = interpreter.get_input_details() + output_details = interpreter.get_output_details() + + + # Prepare input data (replace with your input data) + input_data = Image.fromarray(img) + input_data = input_data.resize((960,608)) #TODO: Extract input size from interpreter + arr_image = np.array(input_data) + img_transform = transforms.Compose([ + transforms.ToTensor() + ]) + input_data = img_transform(input_data).unsqueeze(0) + + + # Set input tensor + interpreter.set_tensor(input_details[0]['index'], input_data) + + + # Run inference + interpreter.invoke() + + + # Get output tensor + output_data = interpreter.get_tensor(output_details[0]['index']) + + + # Process output (interpret predictions) + # Replace this with your post-processing logic + if len(output_data) != 0: + detections = output_data[0] + print(detections[0]) + for x1, y1, w, h, conf, cls1, cls2, cls3, cls4, cls5, cls6 in detections: #TODO: see if you can get an array of classes instead like *args + class_index = np.argmax([cls1,cls2,cls3,cls4,cls5,cls6]) + if conf < CONFIDENCE_THRESHOLD: + continue + else: + pass + print(f'{CLASSES[class_index]} => {conf}') + print(f'Center: ({x1, y1})') + + + x1, y1, x2, y2 = self.__bounding_box_coordinates(x1, y1, w, h) + plot_one_box([x1, y1, x2, y2], arr_image, label=f'{CLASSES[class_index]}', color=COLORS[class_index], line_thickness=2) + else: + print("No Detections Made") + + #TODO: Return post processing image or preprocessed image and the detections Center width and height + return arr_image + + def __bounding_box_coordinates(center_x, center_y, width, height): + half_width = width / 2 + half_height = height / 2 + x1 = center_x - half_width + y1 = center_y - half_height + x2 = center_x + half_width + y2 = center_y + half_height + return x1, y1, x2, y2 \ No newline at end of file From f586e5e039b6d74f333d03d93818b81a6e0dc097 Mon Sep 17 00:00:00 2001 From: Daniel Parra Date: Sun, 24 Mar 2024 17:01:41 -0400 Subject: [PATCH 2/2] Began working on creating the Object detection layer --- layers/ObjectDetectionLayer.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/layers/ObjectDetectionLayer.py b/layers/ObjectDetectionLayer.py index cc7039f..4a0b4f9 100644 --- a/layers/ObjectDetectionLayer.py +++ b/layers/ObjectDetectionLayer.py @@ -15,7 +15,14 @@ from Layer import AnalysisLayer class ObjectDetectionLayer(AnalysisLayer): - def __init__(self, in_size, out_size, path_to_weights, conf, classes_array, colors_array) -> None: + def __init__(self, in_size, out_size, path_to_weights, conf_thres, iou_thres, class_names_array, colors_array) -> None: + # TODO: determine what type of processor to use based on the .type of the path_to_weights string + + # TODO: set the conf and iou threshold values + + # TODO: set the colors array and class_names_array values + + # TODO: set the processor super().__init__(in_size, out_size, "objectDetection") class PTWeightsProcessor():