UPDATE openvino

Vision/tool/utils.py

@@ -1,13 +1,16 @@
 #!/usr/bin/env python
 # -*- coding: UTF-8 -*-
 '''
-@Project ：AutoControlSystem-master 
+@Project ：AutoControlSystem-master
 @File    ：utils.py
-@IDE     ：PyCharm 
+@IDE     ：PyCharm
 @Author  ：hjw
-@Date    ：2024/8/29 15:07 
+@Date    ：2024/8/29 15:07
 '''
 
+import numpy as np
+import cv2
+
 def find_position(Depth_Z, RegionalArea, RegionalArea_Threshold, first_depth=True):
     if first_depth == True:
         sorted_id = sorted(range(len(Depth_Z)), key=lambda k: Depth_Z[k], reverse=False)
@@ -25,6 +28,157 @@ def find_position(Depth_Z, RegionalArea, RegionalArea_Threshold, first_depth=Tru
             if RegionalArea1[i] > RegionalArea_Threshold:
                 return sorted_id[i]
 
+class_names = ['box', 'other']
+
+# Create a list of colors for each class where each color is a tuple of 3 integer values
+rng = np.random.default_rng(3)
+colors = rng.uniform(0, 255, size=(len(class_names), 3))
+
+
+def nms(boxes, scores, iou_threshold):
+    # Sort by score
+    sorted_indices = np.argsort(scores)[::-1]
+
+    keep_boxes = []
+    while sorted_indices.size > 0:
+        # Pick the last box
+        box_id = sorted_indices[0]
+        keep_boxes.append(box_id)
+
+        # Compute IoU of the picked box with the rest
+        ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])
+
+        # Remove boxes with IoU over the threshold
+        keep_indices = np.where(ious < iou_threshold)[0]
+
+        # print(keep_indices.shape, sorted_indices.shape)
+        sorted_indices = sorted_indices[keep_indices + 1]
+
+    return keep_boxes
+
+
+def compute_iou(box, boxes):
+    # Compute xmin, ymin, xmax, ymax for both boxes
+    xmin = np.maximum(box[0], boxes[:, 0])
+    ymin = np.maximum(box[1], boxes[:, 1])
+    xmax = np.minimum(box[2], boxes[:, 2])
+    ymax = np.minimum(box[3], boxes[:, 3])
+
+    # Compute intersection area
+    intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin)
+
+    # Compute union area
+    box_area = (box[2] - box[0]) * (box[3] - box[1])
+    boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
+    union_area = box_area + boxes_area - intersection_area
+
+    # Compute IoU
+    iou = intersection_area / union_area
+
+    return iou
+
+
+def xywh2xyxy(x):
+    # Convert bounding box (x, y, w, h) to bounding box (x1, y1, x2, y2)
+    y = np.copy(x)
+    y[..., 0] = x[..., 0] - x[..., 2] / 2
+    y[..., 1] = x[..., 1] - x[..., 3] / 2
+    y[..., 2] = x[..., 0] + x[..., 2] / 2
+    y[..., 3] = x[..., 1] + x[..., 3] / 2
+    return y
+
+
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
+
+
+def draw_detections(image, boxes, scores, class_ids, mask_alpha=0.3, mask_maps=None):
+    img_height, img_width = image.shape[:2]
+    size = min([img_height, img_width]) * 0.0006
+    text_thickness = int(min([img_height, img_width]) * 0.001)
+
+    mask_img = draw_masks(image, boxes, class_ids, mask_alpha, mask_maps)
+
+    # Draw bounding boxes and labels of detections
+    for box, score, class_id in zip(boxes, scores, class_ids):
+        color = colors[class_id]
+
+        x1, y1, x2, y2 = box.astype(int)
+
+        # Draw rectangle
+        cv2.rectangle(mask_img, (x1, y1), (x2, y2), color, 2)
+
+        label = class_names[class_id]
+        caption = f'{label} {int(score * 100)}%'
+        (tw, th), _ = cv2.getTextSize(text=caption, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+                                      fontScale=size, thickness=text_thickness)
+        th = int(th * 1.2)
+
+        cv2.rectangle(mask_img, (x1, y1),
+                      (x1 + tw, y1 - th), color, -1)
+
+        cv2.putText(mask_img, caption, (x1, y1),
+                    cv2.FONT_HERSHEY_SIMPLEX, size, (255, 255, 255), text_thickness, cv2.LINE_AA)
+
+    return mask_img
+
+
+def draw_masks(image, boxes, class_ids, mask_alpha=0.3, mask_maps=None):
+    mask_img = image.copy()
+
+    # Draw bounding boxes and labels of detections
+    for i, (box, class_id) in enumerate(zip(boxes, class_ids)):
+        color = colors[class_id]
+
+        x1, y1, x2, y2 = box.astype(int)
+
+        # Draw fill mask image
+        if mask_maps is None:
+            cv2.rectangle(mask_img, (x1, y1), (x2, y2), color, -1)
+        else:
+            crop_mask = mask_maps[i][y1:y2, x1:x2, np.newaxis]
+            crop_mask_img = mask_img[y1:y2, x1:x2]
+            crop_mask_img = crop_mask_img * (1 - crop_mask) + crop_mask * color
+            mask_img[y1:y2, x1:x2] = crop_mask_img
+
+    return cv2.addWeighted(mask_img, mask_alpha, image, 1 - mask_alpha, 0)
+
+
+def draw_comparison(img1, img2, name1, name2, fontsize=2.6, text_thickness=3):
+    (tw, th), _ = cv2.getTextSize(text=name1, fontFace=cv2.FONT_HERSHEY_DUPLEX,
+                                  fontScale=fontsize, thickness=text_thickness)
+    x1 = img1.shape[1] // 3
+    y1 = th
+    offset = th // 5
+    cv2.rectangle(img1, (x1 - offset * 2, y1 + offset),
+                  (x1 + tw + offset * 2, y1 - th - offset), (0, 115, 255), -1)
+    cv2.putText(img1, name1,
+                (x1, y1),
+                cv2.FONT_HERSHEY_DUPLEX, fontsize,
+                (255, 255, 255), text_thickness)
+
+    (tw, th), _ = cv2.getTextSize(text=name2, fontFace=cv2.FONT_HERSHEY_DUPLEX,
+                                  fontScale=fontsize, thickness=text_thickness)
+    x1 = img2.shape[1] // 3
+    y1 = th
+    offset = th // 5
+    cv2.rectangle(img2, (x1 - offset * 2, y1 + offset),
+                  (x1 + tw + offset * 2, y1 - th - offset), (94, 23, 235), -1)
+
+    cv2.putText(img2, name2,
+                (x1, y1),
+                cv2.FONT_HERSHEY_DUPLEX, fontsize,
+                (255, 255, 255), text_thickness)
+
+    combined_img = cv2.hconcat([img1, img2])
+    if combined_img.shape[1] > 3840:
+        combined_img = cv2.resize(combined_img, (3840, 2160))
+
+    return combined_img
+
+
+
+