Create new file

lib/core/postprocess.py

@@ -0,0 +1,223 @@
+
+
+import torch
+from lib.utils import is_parallel
+import numpy as np
+np.set_printoptions(threshold=np.inf)
+import cv2
+from sklearn.cluster import DBSCAN
+
+
+def build_targets(cfg, predictions, targets, model):
+    '''
+    predictions
+    [16, 3, 32, 32, 85]
+    [16, 3, 16, 16, 85]
+    [16, 3, 8, 8, 85]
+    torch.tensor(predictions[i].shape)[[3, 2, 3, 2]]
+    [32,32,32,32]
+    [16,16,16,16]
+    [8,8,8,8]
+    targets[3,x,7]
+    t [index, class, x, y, w, h, head_index]
+    '''
+    # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
+    det = model.module.model[model.module.detector_index] if is_parallel(model) \
+        else model.model[model.detector_index]  # Detect() module
+    # print(type(model))
+    # det = model.model[model.detector_index]
+    # print(type(det))
+    na, nt = det.na, targets.shape[0]  # number of anchors, targets
+    tcls, tbox, indices, anch = [], [], [], []
+    gain = torch.ones(7, device=targets.device)  # normalized to gridspace gain
+    ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt)  # same as .repeat_interleave(nt)
+    targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2)  # append anchor indices
+    
+    g = 0.5  # bias
+    off = torch.tensor([[0, 0],
+                        [1, 0], [0, 1], [-1, 0], [0, -1],  # j,k,l,m
+                        # [1, 1], [1, -1], [-1, 1], [-1, -1],  # jk,jm,lk,lm
+                        ], device=targets.device).float() * g  # offsets
+    
+    for i in range(det.nl):
+        anchors = det.anchors[i] #[3,2]
+        gain[2:6] = torch.tensor(predictions[i].shape)[[3, 2, 3, 2]]  # xyxy gain
+        # Match targets to anchors
+        t = targets * gain
+
+        if nt:
+            # Matches
+            r = t[:, :, 4:6] / anchors[:, None]  # wh ratio
+            j = torch.max(r, 1. / r).max(2)[0] < cfg.TRAIN.ANCHOR_THRESHOLD  # compare
+            # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t']  # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
+            t = t[j]  # filter
+
+            # Offsets
+            gxy = t[:, 2:4]  # grid xy
+            gxi = gain[[2, 3]] - gxy  # inverse
+            j, k = ((gxy % 1. < g) & (gxy > 1.)).T
+            l, m = ((gxi % 1. < g) & (gxi > 1.)).T
+            j = torch.stack((torch.ones_like(j), j, k, l, m))
+            t = t.repeat((5, 1, 1))[j]
+            offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
+        else:
+            t = targets[0]
+            offsets = 0
+
+        # Define
+        b, c = t[:, :2].long().T  # image, class
+        gxy = t[:, 2:4]  # grid xy
+        gwh = t[:, 4:6]  # grid wh
+        gij = (gxy - offsets).long()
+        gi, gj = gij.T  # grid xy indices
+
+        # Append
+        a = t[:, 6].long()  # anchor indices
+        indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1)))  # image, anchor, grid indices
+        tbox.append(torch.cat((gxy - gij, gwh), 1))  # box
+        anch.append(anchors[a])  # anchors
+        tcls.append(c)  # class
+
+    return tcls, tbox, indices, anch
+
+def morphological_process(image, kernel_size=5, func_type=cv2.MORPH_CLOSE):
+    """
+    morphological process to fill the hole in the binary segmentation result
+    :param image:
+    :param kernel_size:
+    :return:
+    """
+    if len(image.shape) == 3:
+        raise ValueError('Binary segmentation result image should be a single channel image')
+
+    if image.dtype is not np.uint8:
+        image = np.array(image, np.uint8)
+
+    kernel = cv2.getStructuringElement(shape=cv2.MORPH_ELLIPSE, ksize=(kernel_size, kernel_size))
+
+    # close operation fille hole
+    closing = cv2.morphologyEx(image, func_type, kernel, iterations=1)
+
+    return closing
+
+def connect_components_analysis(image):
+    """
+    connect components analysis to remove the small components
+    :param image:
+    :return:
+    """
+    if len(image.shape) == 3:
+        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    else:
+        gray_image = image
+    # print(gray_image.dtype)
+    return cv2.connectedComponentsWithStats(gray_image, connectivity=8, ltype=cv2.CV_32S)
+
+def if_y(samples_x):
+    for sample_x in samples_x:
+        if len(sample_x):
+            # if len(sample_x) != (sample_x[-1] - sample_x[0] + 1) or sample_x[-1] == sample_x[0]:
+            if sample_x[-1] == sample_x[0]:
+                return False
+    return True
+    
+def fitlane(mask, sel_labels, labels, stats):
+    H, W = mask.shape
+    for label_group in sel_labels:
+        states = [stats[k] for k in label_group]
+        x, y, w, h, _ = states[0]
+        # if len(label_group) > 1:
+        #     print('in')
+        #     for m in range(len(label_group)-1):
+        #         labels[labels == label_group[m+1]] = label_group[0]
+        t = label_group[0]
+        # samples_y = np.linspace(y, H-1, 30)
+        # else:
+        samples_y = np.linspace(y, y+h-1, 30)
+        
+        samples_x = [np.where(labels[int(sample_y)]==t)[0] for sample_y in samples_y]
+
+        if if_y(samples_x):
+            samples_x = [int(np.mean(sample_x)) if len(sample_x) else -1 for sample_x in samples_x]
+            samples_x = np.array(samples_x)
+            samples_y = np.array(samples_y)
+            samples_y = samples_y[samples_x != -1]
+            samples_x = samples_x[samples_x != -1]
+            func = np.polyfit(samples_y, samples_x, 2)
+            x_limits = np.polyval(func, H-1)
+            # if (y_max + h - 1) >= 720:
+            if x_limits < 0 or x_limits > W:
+            # if (y_max + h - 1) > 720:
+                # draw_y = np.linspace(y, 720-1, 720-y)
+                draw_y = np.linspace(y, y+h-1, h)
+            else:
+                # draw_y = np.linspace(y, y+h-1, y+h-y)
+                draw_y = np.linspace(y, H-1, H-y)
+            draw_x = np.polyval(func, draw_y)
+            # draw_y = draw_y[draw_x < W]
+            # draw_x = draw_x[draw_x < W]
+            draw_points = (np.asarray([draw_x, draw_y]).T).astype(np.int32)
+            cv2.polylines(mask, [draw_points], False, 1, thickness=15)
+        else:
+            # if ( + w - 1) >= 1280:
+            samples_x = np.linspace(x, W-1, 30)
+            # else:
+            #     samples_x = np.linspace(x, x_max+w-1, 30)
+            samples_y = [np.where(labels[:, int(sample_x)]==t)[0] for sample_x in samples_x]
+            samples_y = [int(np.mean(sample_y)) if len(sample_y) else -1 for sample_y in samples_y]
+            samples_x = np.array(samples_x)
+            samples_y = np.array(samples_y)
+            samples_x = samples_x[samples_y != -1]
+            samples_y = samples_y[samples_y != -1]
+            try:
+                func = np.polyfit(samples_x, samples_y, 2)
+            except:
+                pass
+            # y_limits = np.polyval(func, 0)
+            # if y_limits > 720 or y_limits < 0:
+            # if (x + w - 1) >= 1280:
+            #     draw_x = np.linspace(x, 1280-1, 1280-x)
+            # else:
+            y_limits = np.polyval(func, 0)
+            if y_limits >= H or y_limits < 0:
+                draw_x = np.linspace(x, x+w-1, w+x-x)
+            else:
+                y_limits = np.polyval(func, W-1)
+                if y_limits >= H or y_limits < 0:
+                    draw_x = np.linspace(x, x+w-1, w+x-x)
+                # if x+w-1 < 640:
+                #     draw_x = np.linspace(0, x+w-1, w+x-x)
+                else:
+                    draw_x = np.linspace(x, W-1, W-x)
+            draw_y = np.polyval(func, draw_x)
+            draw_points = (np.asarray([draw_x, draw_y]).T).astype(np.int32)
+            cv2.polylines(mask, [draw_points], False, 1, thickness=15)
+    return mask
+
+def connect_lane(image, shadow_height=0):
+    if len(image.shape) == 3:
+        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    else:
+        gray_image = image
+    if shadow_height:
+        image[:shadow_height] = 0
+    mask = np.zeros((image.shape[0], image.shape[1]), np.uint8)
+    
+    num_labels, labels, stats, centers = cv2.connectedComponentsWithStats(gray_image, connectivity=8, ltype=cv2.CV_32S)
+    # ratios = []
+    selected_label = []
+    
+    for t in range(1, num_labels, 1):
+        _, _, _, _, area = stats[t]
+        if area > 400:
+            selected_label.append(t)
+    if len(selected_label) == 0:
+        return mask
+    else:
+        split_labels = [[label,] for label in selected_label]
+        mask_post = fitlane(mask, split_labels, labels, stats)
+        return mask_post
+
+
+
+