Upload tf_post_processing.py

tf_post_processing.py

@@ -0,0 +1,233 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Mon Sep  4 16:03:42 2023
+
+@author: SABARI
+"""
+import time
+import tensorflow as tf
+import numpy as np
+#from lsnms import nms, wbc
+
+def box_iou(box1, box2, eps=1e-7):
+    """
+    Calculate intersection-over-union (IoU) of boxes.
+    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
+    Args:
+    box1 (tf.Tensor): A tensor of shape (N, 4) representing N bounding boxes.
+    box2 (tf.Tensor): A tensor of shape (M, 4) representing M bounding boxes.
+    eps (float, optional): A small value to avoid division by zero. Defaults to 1e-7.
+    Returns:
+    (tf.Tensor): An NxM tensor containing the pairwise IoU values for every element in box1 and box2.
+    """
+    
+    a1, a2 = tf.split(box1, 2, axis=1)
+    b1, b2 = tf.split(box2, 2, axis=1)
+    
+    inter = tf.reduce_prod(tf.maximum(tf.minimum(a2, b2) - tf.maximum(a1, b1), 0), axis=1)
+    
+    return inter / (tf.reduce_prod(a2 - a1, axis=1) + tf.reduce_prod(b2 - b1, axis=1) - inter + eps)
+
+
+def xywh2xyxy(x):
+    """
+    Convert bounding box coordinates from (x, y, width, height) format to (x1, y1, x2, y2) format where (x1, y1) is the
+    top-left corner and (x2, y2) is the bottom-right corner.
+    Args:
+    x (tf.Tensor): The input bounding box coordinates in (x, y, width, height) format.
+    Returns:
+    y (tf.Tensor): The bounding box coordinates in (x1, y1, x2, y2) format.
+    """
+    # Assuming x is a NumPy array
+    y = np.copy(x)
+    y[..., 0] = x[..., 0] - x[..., 2] / 2  # top left x
+    y[..., 1] = x[..., 1] - x[..., 3] / 2  # top left y
+    y[..., 2] = x[..., 0] + x[..., 2] / 2  # bottom right x
+    y[..., 3] = x[..., 1] + x[..., 3] / 2  # bottom right y
+    return y
+
+
+def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, agnostic=False,
+                        multi_label=False, max_det=300, nc=0,  # number of classes (optional)
+                        max_time_img=0.05,
+                        max_nms=100,
+                        max_wh=7680):
+    """
+    Perform non-maximum suppression (NMS) on a set of boxes, with support for masks and multiple labels per box.
+    Arguments:
+    prediction (tf.Tensor): A tensor of shape (batch_size, num_classes + 4 + num_masks, num_boxes)
+    containing the predicted boxes, classes, and masks. The tensor should be in the format
+    output by a model, such as YOLO.
+    conf_thres (float): The confidence threshold below which boxes will be filtered out.
+    Valid values are between 0.0 and 1.0.
+    iou_thres (float): The IoU threshold below which boxes will be filtered out during NMS.
+    Valid values are between 0.0 and 1.0.
+    agnostic (bool): If True, the model is agnostic to the number of classes, and all
+    classes will be considered as one.
+    multi_label (bool): If True, each box may have multiple labels.
+    max_det (int): The maximum number of boxes to keep after NMS.
+    nc (int): (optional) The number of classes output by the model. Any indices after this will be considered masks.
+    max_time_img (float): The maximum time (seconds) for processing one image.
+    max_nms (int): The maximum number of boxes into tf.image.combined_non_max_suppression().
+    max_wh (int): The maximum box width and height in pixels
+    Returns:
+    (List[tf.Tensor]): A list of length batch_size, where each element is a tensor of
+    shape (num_boxes, 6 + num_masks) containing the kept boxes, with columns
+    (x1, y1, x2, y2, confidence, class, mask1, mask2, ...).
+    """
+
+    # Checks
+    assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
+    assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
+    if isinstance(prediction, (list, tuple)):  # YOLOv8 model in validation model, output = (inference_out, loss_out)
+        prediction = prediction[0]  # select only inference output
+
+    bs = np.shape(prediction)[0]  # batch size
+    nc = nc or (np.shape(prediction)[1] - 4)  # number of classes
+    nm = np.shape(prediction)[1] - nc - 4
+    mi = 4 + nc  # mask start index
+    #xc = tf.math.reduce_any(prediction[:, 4:mi] > conf_thres, axis=1)  # candidates
+    xc = np.amax(prediction[:, 4:mi], axis=1) > conf_thres
+
+    # Settings
+    # min_wh = 2  # (pixels) minimum box width and height
+    time_limit = 0.5 + max_time_img * tf.cast(bs, tf.float32)  # seconds to quit after
+    multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
+
+    t = time.time()
+    output = [np.zeros((0, 6 + nm))] * bs
+    for xi, x in enumerate(prediction):  # image index, image inference
+        # Apply constraints
+        # x = tf.where(tf.math.logical_or(x[:, 2:4] < min_wh, x[:, 2:4] > max_wh), tf.constant(0, dtype=tf.float32), x)  # width-height
+        #x = tf.boolean_mask(x, xc[xi])
+        #x = x.transpose(0, -1)[xc[xi]]  # confidence
+        # Assuming x, xc, and xi are NumPy arrays
+        x = np.transpose(x)
+
+        #x = x.transpose()[:, xc[xi]]
+        x = x[xc[xi]]
+
+        # If none remain process next image
+        if np.shape(x)[0] == 0:
+            continue
+
+        # Detections matrix nx6 (xyxy, conf, cls)
+        #box, cls, mask = tf.split(x, [4, nc, nm], axis=1)
+        
+        # Assuming x is a NumPy array
+        box = x[:, :4]
+        cls = x[:, 4:4 + nc]
+        mask = x[:, 4 + nc:]
+        box = xywh2xyxy(box)  # center_x, center_y, width, height) to (x1, y1, x2, y2)
+
+        # Assuming cls is a NumPy array
+        if multi_label:
+            i, j = np.where(cls > conf_thres)
+            x = np.concatenate([box[i], np.expand_dims(cls[i, j], axis=-1), np.expand_dims(j, axis=-1).astype(np.float32), mask[i]], axis=1)
+        else:
+            conf = np.max(cls, axis=1)
+            j = np.argmax(cls, axis=1)
+            keep = np.where(conf > conf_thres)[0]
+            x = np.concatenate([box[keep], np.expand_dims(conf[keep], axis=-1), np.expand_dims(j[keep], axis=-1).astype(np.float32), mask[keep]], axis=1)
+
+        # Check shape
+        n = np.shape(x)[0]  # number of boxes
+        if n == 0:  # no boxes
+            continue
+        #x = x[tf.argsort(x[:, 4], direction='DESCENDING')[:max_nms]]  # sort by confidence and remove excess boxes
+        sorted_indices = np.argsort(x[:, 4])[::-1]  # Sort indices in descending order of confidence
+        x = x[sorted_indices[:max_nms]]  # Keep the top max_nms boxes
+        
+        # Batched NMS
+        c = x[:, 5:6] * (0.0 if agnostic else tf.cast(max_wh, tf.float32))  # classes
+        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
+        i = tf.image.non_max_suppression(boxes, scores, max_nms, iou_threshold=iou_thres)  # NMS
+        i = i.numpy()
+        i = i[:max_det]  # limit detections
+
+        output[xi] =  x[i,:]
+        
+        if (time.time() - t) > time_limit:
+            break  # time limit exceeded
+
+    return output
+
+import numpy as np
+
+def optimized_object_detection(prediction, conf_thres=0.25, iou_thres=0.45, agnostic=False,
+                                multi_label=False, max_det=300, nc=0, max_time_img=0.05,
+                                max_nms=100, max_wh=7680):
+    
+    assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
+    assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
+    
+    if isinstance(prediction, (list, tuple)):
+        prediction = prediction[0]
+    
+    bs, _, _ = prediction.shape  # Get batch size and dimensions
+    
+    if nc == 0:
+        nc = prediction.shape[1] - 4
+    
+    nm = prediction.shape[1] - nc - 4
+    mi = 4 + nc
+    
+    xc = np.amax(prediction[:, 4:mi], axis=1) > conf_thres
+    
+    time_limit = 0.5 + max_time_img * bs
+    
+    multi_label &= nc > 1
+    
+    t = time.time()
+    output = [np.zeros((0, 6 + nm))] * bs
+    
+    for xi, x in enumerate(prediction):
+        x = np.transpose(x)
+        x = x[xc[xi]]
+        
+        if np.shape(x)[0] == 0:
+            continue
+
+        box = x[:, :4]
+        cls = x[:, 4:4 + nc]
+        mask = x[:, 4 + nc:]
+        box = xywh2xyxy(box)
+
+        if multi_label:
+            i, j = np.where(cls > conf_thres)
+            x = np.concatenate([box[i], np.expand_dims(cls[i, j], axis=-1), np.expand_dims(j, axis=-1).astype(np.float32), mask[i]], axis=1)
+        else:
+            conf = np.max(cls, axis=1)
+            j = np.argmax(cls, axis=1)
+            keep = np.where(conf > conf_thres)[0]
+            x = np.concatenate([box[keep], np.expand_dims(conf[keep], axis=-1), np.expand_dims(j[keep], axis=-1).astype(np.float32), mask[keep]], axis=1)
+
+        n = np.shape(x)[0]
+        if n == 0:
+            continue
+        
+        sorted_indices = np.argsort(x[:, 4])[::-1]
+        x = x[sorted_indices[:max_nms]]
+        
+        c = x[:, 5:6] * (0.0 if agnostic else max_wh)
+        boxes, scores = x[:, :4] + c, x[:, 4]
+        i = tf.image.non_max_suppression(boxes, scores, max_nms, iou_threshold=iou_thres)
+        
+        #keep = nms(boxes, scores, iou_threshold=iou_thres)
+
+        i = i.numpy()
+        i = i[:max_det]
+        
+        output[xi] = x[keep,:]
+        
+        if (time.time() - t) > time_limit:
+            break
+    
+    return output
+
+
+#output_numpy = np.load(r"D:\object_face_person_detection\yolov8_tf_results\gustavo-alves-YOXSC4zRcxw-unsplash.npy")
+
+#detections = non_max_suppression(output_numpy, conf_thres=0.4, iou_thres=0.4)[0]
+
+#print(detections)