Initial commit

TheDistanceAssessor.py

@@ -0,0 +1,905 @@
+import cv2
+import os
+import time
+import json
+import numpy as np
+import matplotlib.pyplot as plt
+from sklearn.linear_model import LinearRegression
+import matplotlib.path as mplPath
+import matplotlib.patches as patches
+from ultralyticsplus import YOLO
+from scripts.test_filtered_cls import load, load_model, process
+
+PATH_jpgs = 'assets/rs19val/jpgs/test'
+PATH_model_seg = 'assets/models_pretrained/segformer/SegFormer_B3_1024_finetuned.pth'
+PATH_model_det = 'assets/models_pretrained/ultralyticsplus/yolov8s'
+PATH_base = 'assets/pilsen_railway_dataset/'
+eda_path = "assets/pilsen_railway_dataset/eda_table.table.json"
+data_json = json.load(open(eda_path, 'r'))
+
+def load_yolo(PATH_model):
+        model = YOLO(PATH_model)
+
+        model.overrides['conf'] = 0.25  # NMS confidence threshold
+        model.overrides['iou'] = 0.45  # NMS IoU threshold
+        model.overrides['agnostic_nms'] = False  # NMS class-agnostic
+        model.overrides['max_det'] = 1000  # maximum number of detections per image
+        return model
+
+def find_extreme_y_values(arr, values=[0, 6]):
+        """
+        Optimized function to find the lowest and highest y-values (row indices) in a 2D array where 0 or 6 appears.
+        
+        Parameters:
+        - arr: The input 2D NumPy array.
+        - values: The values to search for (default is [0, 6]).
+        
+        Returns:
+        A tuple (lowest_y, highest_y) representing the lowest and highest y-values. If values are not found, returns None.
+        """
+        mask = np.isin(arr, values)
+        rows_with_values = np.any(mask, axis=1)
+        
+        y_indices = np.nonzero(rows_with_values)[0]  # Directly finding non-zero (True) indices
+        
+        if y_indices.size == 0:
+                return None, None  # Early return if values not found
+        
+        return y_indices[0], y_indices[-1]
+
+def find_nearest_pairs(arr1, arr2):
+        # Convert lists to numpy arrays for vectorized operations
+        arr1_np = np.array(arr1)
+        arr2_np = np.array(arr2)
+        
+        # Determine which array is shorter
+        if len(arr1_np) < len(arr2_np):
+                base_array, compare_array = arr1_np, arr2_np
+        else:
+                base_array, compare_array = arr2_np, arr1_np
+
+        paired_base = []
+        paired_compare = []
+
+        # Mask to keep track of paired elements
+        paired_mask = np.zeros(len(compare_array), dtype=bool)
+
+        for item in base_array:
+                # Calculate distances from the current item to all items in the compare_array
+                distances = np.linalg.norm(compare_array - item, axis=1)
+                nearest_index = np.argmin(distances)
+                paired_base.append(item)
+                paired_compare.append(compare_array[nearest_index])
+                # Mark the paired element to exclude it from further pairing
+                paired_mask[nearest_index] = True
+
+                # Check if all elements from the compare_array have been paired
+                if paired_mask.all():
+                        break
+
+        paired_base = np.array(paired_base)
+        paired_compare = compare_array[paired_mask]
+
+        return (paired_base, paired_compare) if len(arr1_np) < len(arr2_np) else (paired_compare, paired_base)
+
+def filter_crossings(image, edges_dict):
+        filtered_edges = {}
+        for key, values in edges_dict.items():
+                merged = [values[0]]
+                for start, end in values[1:]:
+                        if start - merged[-1][1] < 50:
+                                
+                                key_up = max([0, key-10])
+                                key_down = min([image.shape[0]-1, key+10])
+                                if key_up == 0:
+                                        key_up = key+20
+                                if key_down == image.shape[0]-1:
+                                        key_down = key-20
+                                
+                                edges_to_test_slope1 = robust_edges(image, [key_up], values=[0, 6], min_width=19)
+                                edges_to_test_slope2 = robust_edges(image, [key_down], values=[0, 6], min_width=19)
+                                
+                                values1, edges_to_test_slope1 = find_nearest_pairs(values, edges_to_test_slope1)
+                                values2, edges_to_test_slope2 = find_nearest_pairs(values, edges_to_test_slope2)
+                                
+                                differences_y = []
+                                for i, value in enumerate(values1):
+                                        if start in value:
+                                                idx = list(value).index(start)
+                                                try:
+                                                        differences_y.append(abs(start-edges_to_test_slope1[i][idx]))
+                                                except:
+                                                        pass
+                                        if merged[-1][1] in value:
+                                                idx = list(value).index(merged[-1][1])
+                                                try:
+                                                        differences_y.append(abs(merged[-1][1]-edges_to_test_slope1[i][idx]))
+                                                except:
+                                                        pass
+                                for i, value in enumerate(values2):
+                                        if start in value:
+                                                idx = list(value).index(start)
+                                                try:
+                                                        differences_y.append(abs(start-edges_to_test_slope2[i][idx]))
+                                                except:
+                                                        pass
+                                        if merged[-1][1] in value:
+                                                idx = list(value).index(merged[-1][1])
+                                                try:
+                                                        differences_y.append(abs(merged[-1][1]-edges_to_test_slope2[i][idx]))
+                                                except:
+                                                        pass
+                                
+                                if any(element > 30 for element in differences_y):
+                                        merged[-1] = (merged[-1][0], end)
+                                else:
+                                        merged.append((start, end))
+                        else:
+                                merged.append((start, end))
+                filtered_edges[key] = merged
+                
+        return filtered_edges
+
+def robust_edges(image, y_levels, values=[0, 6], min_width=19):
+        
+        for y in y_levels:
+                row = image[y, :]
+                mask = np.isin(row, values).astype(int)
+                padded_mask = np.pad(mask, (1, 1), 'constant', constant_values=0)
+                diff = np.diff(padded_mask)
+                starts = np.where(diff == 1)[0]
+                ends = np.where(diff == -1)[0] - 1
+
+                # Filter sequences based on the minimum width criteria
+                filtered_edges = [(start, end) for start, end in zip(starts, ends) if end - start + 1 >= min_width]
+                filtered_edges = [(start, end) for start, end in filtered_edges if 0 not in (start, end) and 1919 not in (start, end)]
+        
+        return filtered_edges
+
+def find_edges(image, y_levels, values=[0, 6], min_width=19):
+        """
+        Find start and end positions of continuous sequences of specified values at given y-levels in a 2D array,
+        filtering for sequences that meet or exceed a specified minimum width.
+
+        Parameters:
+        - arr: 2D NumPy array to search within.
+        - y_levels: List of y-levels (row indices) to examine.
+        - values: Values to search for (default is [0, 6]).
+        - min_width: Minimum width of sequences to be included in the results.
+
+        Returns:
+        A dict with y-levels as keys and lists of (start, end) tuples for each sequence found in that row that meets the width criteria.
+        """
+        edges_dict = {}
+        for y in y_levels:
+                row = image[y, :]
+                mask = np.isin(row, values).astype(int)
+                padded_mask = np.pad(mask, (1, 1), 'constant', constant_values=0)
+                diff = np.diff(padded_mask)
+                starts = np.where(diff == 1)[0]
+                ends = np.where(diff == -1)[0] - 1
+
+                # Filter sequences based on the minimum width criteria
+                filtered_edges = [(start, end) for start, end in zip(starts, ends) if end - start + 1 >= min_width]
+                filtered_edges = [(start, end) for start, end in filtered_edges if 0 not in (start, end) and 1919 not in (start, end)]
+                
+                edges_with_guard_rails = []
+                for edge in filtered_edges:
+                        cutout_left = image[y,edge[0]-50:edge[0]][::-1]
+                        cutout_right = image[y,edge[1]:edge[1]+50]
+                        
+                        not_ones = np.where(cutout_left != 1)[0]
+                        if len(not_ones) > 0 and not_ones[0] > 0:
+                                last_one_index = not_ones[0] - 1
+                                edge = (edge[0] - last_one_index,) + edge[1:]
+                        else:
+                                last_one_index = None if len(not_ones) == 0 else not_ones[-1] - 1
+                        
+                        not_ones = np.where(cutout_right != 1)[0]
+                        if len(not_ones) > 0 and not_ones[0] > 0:
+                                last_one_index = not_ones[0] - 1
+                                edge = (edge[0], edge[1] - last_one_index) + edge[2:]
+                        else:
+                                last_one_index = None if len(not_ones) == 0 else not_ones[-1] - 1
+                        
+                        edges_with_guard_rails.append(edge)
+
+                edges_dict[y] = edges_with_guard_rails
+        
+        edges_dict = {k: v for k, v in edges_dict.items() if v}
+        
+        edges_dict = filter_crossings(image, edges_dict)
+        
+        return edges_dict
+
+def find_rails(arr, y_levels, values=[9, 10], min_width=5):
+        edges_all = []
+        for y in y_levels:
+                row = arr[y, :]
+                mask = np.isin(row, values).astype(int)
+                padded_mask = np.pad(mask, (1, 1), 'constant', constant_values=0)
+                diff = np.diff(padded_mask)
+                starts = np.where(diff == 1)[0]
+                ends = np.where(diff == -1)[0] - 1
+
+                # Filter sequences based on the minimum width criteria
+                filtered_edges = [(start, end) for start, end in zip(starts, ends) if end - start + 1 >= min_width]
+                filtered_edges = [(start, end) for start, end in filtered_edges if 0 not in (start, end) and 1919 not in (start, end)]
+                edges_all = filtered_edges
+        
+        return edges_all
+
+def mark_edges(arr, edges_dict, mark_value=30):
+        """
+        Marks a 5x5 zone around the edges found in the array with a specific value.
+
+        Parameters:
+        - arr: The original 2D NumPy array.
+        - edges_dict: A dictionary with y-levels as keys and lists of (start, end) tuples for edges.
+        - mark_value: The value used to mark the edges.
+
+        Returns:
+        The modified array with marked zones.
+        """
+        marked_arr = np.copy(arr)  # Create a copy of the array to avoid modifying the original
+        offset = 2  # To mark a 5x5 area, we go 2 pixels in each direction from the center
+
+        for y, edges in edges_dict.items():
+                for start, end in edges:
+                        # Mark a 5x5 zone around the start and end positions
+                        for dy in range(-offset, offset + 1):
+                                for dx in range(-offset, offset + 1):
+                                        # Check array bounds before marking
+                                        if 0 <= y + dy < marked_arr.shape[0] and 0 <= start + dx < marked_arr.shape[1]:
+                                                marked_arr[y + dy, start + dx] = mark_value
+                                        if 0 <= y + dy < marked_arr.shape[0] and 0 <= end + dx < marked_arr.shape[1]:
+                                                marked_arr[y + dy, end + dx] = mark_value
+
+        return marked_arr
+
+def find_rail_sides(img, edges_dict):
+        left_border = []
+        right_border = []
+        for y,xs in edges_dict.items():
+                rails = find_rails(img, [y], values=[9,10], min_width=5)
+                left_border_actual = [min(xs)[0],y]
+                right_border_actual = [max(xs)[1],y]
+                
+                for zone in rails:
+                        if abs(zone[1]-left_border_actual[0]) < y*0.04: # dynamic treshold
+                                left_border_actual[0] = zone[0]
+                        if abs(zone[0]-right_border_actual[0]) < y*0.04:
+                                right_border_actual[0] = zone[1]
+                
+                left_border.append(left_border_actual)
+                right_border.append(right_border_actual)
+
+        # removing detected uncontioussness
+        left_border, flags_l, _ = robust_rail_sides(left_border) # filter outliers
+        right_border, flags_r, _ = robust_rail_sides(right_border)
+        
+        return left_border, right_border, flags_l, flags_r
+
+def robust_rail_sides(border, threshold=7):
+        border = np.array(border)
+        if border.size > 0:
+                # delete borders found on the bottom side of the image
+                border = border[border[:, 1] != 1079]
+                
+                steps_x = np.diff(border[:, 0])
+                median_step = np.median(np.abs(steps_x))
+                
+                threshold_step = np.abs(threshold*np.abs(median_step))
+                treshold_overcommings = abs(steps_x) > abs(threshold_step)
+                
+                flags = []
+                
+                if True not in treshold_overcommings:
+                        return border, flags, []
+                else:
+                        overcommings_indices = [i for i, element in enumerate(treshold_overcommings) if element == True]
+                        if overcommings_indices and np.all(np.diff(overcommings_indices) == 1):
+                                overcommings_indices = [overcommings_indices[0]]
+                        
+                        filtered_border = border
+                        
+                        previously_deleted = []
+                        for i in overcommings_indices:
+                                for item in previously_deleted:
+                                        if item[0] < i:
+                                                i -= item[1]
+                                first_part = filtered_border[:i+1]
+                                second_part = filtered_border[i+1:]
+                                if len(second_part)<2:
+                                        filtered_border = first_part
+                                        previously_deleted.append([i,len(second_part)])
+                                elif len(first_part)<2:
+                                        filtered_border = second_part
+                                        previously_deleted.append([i,len(first_part)])
+                                else:
+                                        first_b, _, deleted_first = robust_rail_sides(first_part)
+                                        second_b, _, _ = robust_rail_sides(second_part)
+                                        filtered_border = np.concatenate((first_b,second_b), axis=0)
+                                        
+                                        if deleted_first:
+                                                for deleted_item in deleted_first:
+                                                        if deleted_item[0]<=i:
+                                                                i -= deleted_item[1]
+                                                
+                                        flags.append(i)
+                        return filtered_border, flags, previously_deleted
+        else:
+                return border, [], []
+
+def find_dist_from_edges(id_map, image, edges_dict, left_border, right_border, real_life_width_mm, real_life_target_mm, mark_value=30):
+        """
+        Mark regions representing a real-life distance (e.g., 2 meters) to the left and right from the furthest edges.
+        
+        Parameters:
+        - arr: 2D NumPy array representing the id_map.
+        - edges_dict: Dictionary with y-levels as keys and lists of (start, end) tuples for edges.
+        - real_life_width_mm: The real-world width in millimeters that the average sequence width represents.
+        - real_life_target_mm: The real-world distance in millimeters to mark from the edges.
+        
+        Returns:
+        - A NumPy array with the marked regions.
+        """
+        # Calculate the rail widths
+        diffs_widths = {k: sum(e-s for s, e in v) / len(v) for k, v in edges_dict.items() if v}
+        diffs_width = {k: max(e-s for s, e in v) for k, v in edges_dict.items() if v}
+
+        # Pixel to mm scale factor
+        scale_factors = {k: real_life_width_mm / v for k, v in diffs_width.items()}
+        # Converting the real-life target distance to pixels
+        target_distances_px = {k: int(real_life_target_mm / v) for k, v in scale_factors.items()}
+        
+        # Mark the regions representing the target distance to the left and right from the furthest edges
+        end_points_left = {}
+        region_levels_left = []
+        for point in left_border:
+                min_edge = point[0]
+                
+                # Ensure we stay within the image bounds
+                #left_mark_start = max(0, min_edge - int(target_distances_px[point[1]]))
+                left_mark_start = min_edge - int(target_distances_px[point[1]])
+                end_points_left[point[1]] = left_mark_start
+                
+                # Left region points
+                if left_mark_start < min_edge:
+                        y_values = np.arange(left_mark_start, min_edge)
+                        x_values = np.full_like(y_values, point[1])
+                        region_line = np.column_stack((x_values, y_values))
+                        region_levels_left.append(region_line)
+                        
+        end_points_right = {}
+        region_levels_right = []
+        for point in right_border:
+                max_edge = point[0]
+                
+                # Ensure we stay within the image bounds
+                right_mark_end = min(id_map.shape[1], max_edge + int(target_distances_px[point[1]]))
+                if right_mark_end != id_map.shape[1]:
+                        end_points_right[point[1]] = right_mark_end
+
+                # Right region points
+                if max_edge < right_mark_end:
+                        y_values = np.arange(max_edge, right_mark_end)
+                        x_values = np.full_like(y_values, point[1])
+                        region_line = np.column_stack((x_values, y_values))
+                        region_levels_right.append(region_line)
+
+        return id_map, end_points_left, end_points_right, region_levels_left, region_levels_right
+
+def bresenham_line(x0, y0, x1, y1):
+        """
+        Generate the coordinates of a line from (x0, y0) to (x1, y1) using Bresenham's algorithm.
+        """
+        line = []
+        dx = abs(x1 - x0)
+        dy = -abs(y1 - y0)
+        sx = 1 if x0 < x1 else -1
+        sy = 1 if y0 < y1 else -1
+        err = dx + dy  # error value e_xy
+
+        while True:
+                line.append((x0, y0))  # Add the current point to the line
+                if x0 == x1 and y0 == y1:
+                        break
+                e2 = 2 * err
+                if e2 >= dy:  # e_xy+e_x > 0
+                        err += dy
+                        x0 += sx
+                if e2 <= dx:  # e_xy+e_y < 0
+                        err += dx
+                        y0 += sy
+
+        return line
+
+def interpolate_end_points(end_points_dict, flags):
+        line_arr = []
+        ys = list(end_points_dict.keys())
+        xs = list(end_points_dict.values())
+        
+        if flags and len(flags) == 1:
+                pass
+        elif flags and np.all(np.diff(flags) == 1):
+                flags = [flags[0]]
+        
+        for i in range(0, len(ys) - 1):
+                if i in flags:
+                        continue
+                y1, y2 = ys[i], ys[i + 1]
+                x1, x2 = xs[i], xs[i + 1]
+                line = np.array(bresenham_line(x1, y1, x2, y2))
+                if np.any(line[:, 0] < 0):
+                        line = line[line[:, 0] > 0]
+                line_arr = line_arr + list(line)
+        
+        return line_arr
+
+def extrapolate_line(pixels, image, min_y=None, extr_pixels=10):
+        """
+        Extrapolate a line based on the last segment using linear regression.
+        
+        Parameters:
+        - pixels: List of (x, y) tuples representing line pixel coordinates.
+        - image: 2D numpy array representing the image.
+        - min_y: Minimum y-value to extrapolate to (optional).
+        
+        Returns:
+        - A list of new extrapolated (x, y) pixel coordinates.
+        """
+        if len(pixels) < extr_pixels:
+                print("Not enough pixels to perform extrapolation.")
+                return []
+
+        recent_pixels = np.array(pixels[-extr_pixels:])
+        
+        X = recent_pixels[:, 0].reshape(-1, 1)  # Reshape for sklearn
+        y = recent_pixels[:, 1]
+        
+        model = LinearRegression()
+        model.fit(X, y)
+        
+        slope = model.coef_[0]
+        intercept = model.intercept_
+
+        extrapolate = lambda x: slope * x + intercept
+        
+        # Calculate direction based on last two pixels
+        dx, dy = 0, 0  # Default values
+        
+        x_diffs = []
+        y_diffs = []
+        for i in range(1,extr_pixels-1):
+                x_diffs.append(pixels[-i][0] - pixels[-(i+1)][0])
+                y_diffs.append(pixels[-i][1] - pixels[-(i+1)][1])
+                
+        x_diff = x_diffs[np.argmax(np.abs(x_diffs))]
+        y_diff = y_diffs[np.argmax(np.abs(y_diffs))]
+        
+        if abs(int(x_diff)) >= abs(int(y_diff)):
+                dx = 1 if x_diff >= 0 else -1
+        else:
+                dy = 1 if y_diff >= 0 else -1
+
+        last_pixel = pixels[-1]
+        new_pixels = []
+        x, y = last_pixel
+
+        min_y = min_y if min_y is not None else image.shape[0] - 1
+        
+        while 0 <= x < image.shape[1] and min_y <= y < image.shape[0]:
+                if dx != 0:  # Horizontal or diagonal movement
+                        x += dx
+                        y = int(extrapolate(x))
+                elif dy != 0:  # Vertical movement
+                        y += dy
+                        # For vertical lines, approximate x based on the last known value
+                        x = int(x)
+                        
+                if 0 <= y < image.shape[0] and 0 <= x < image.shape[1]:
+                        new_pixels.append((x, y))
+                else:
+                        break
+
+        return new_pixels
+
+def extrapolate_borders(dist_marked_id_map, border_l, border_r, lowest_y):
+        
+        #border_extrapolation_l1 = extrapolate_line(border_l, dist_marked_id_map, lowest_y)
+        border_extrapolation_l2 = extrapolate_line(border_l[::-1], dist_marked_id_map, lowest_y)
+        
+        #border_extrapolation_r1 = extrapolate_line(border_r, dist_marked_id_map, lowest_y)
+        border_extrapolation_r2 = extrapolate_line(border_r[::-1], dist_marked_id_map, lowest_y)
+        
+        #border_l = border_extrapolation_l2[::-1] + border_l + border_extrapolation_l1
+        #border_r = border_extrapolation_r2[::-1] + border_r + border_extrapolation_r1
+        
+        border_l = border_extrapolation_l2[::-1] + border_l
+        border_r = border_extrapolation_r2[::-1] + border_r
+        
+        return border_l, border_r
+
+def find_zone_border(id_map, image, edges, irl_width_mm=1435, irl_target_mm=1000, lowest_y = 0):
+        
+        left_border, right_border, flags_l, flags_r = find_rail_sides(id_map, edges)
+        
+        dist_marked_id_map, end_points_left, end_points_right, left_region, right_region = find_dist_from_edges(id_map, image, edges, left_border, right_border, irl_width_mm, irl_target_mm)
+        
+        border_l = interpolate_end_points(end_points_left, flags_l)
+        border_r = interpolate_end_points(end_points_right, flags_r)
+        
+        border_l, border_r = extrapolate_borders(dist_marked_id_map, border_l, border_r, lowest_y)
+        
+        return [border_l, border_r],[left_region, right_region]
+
+def get_clues(segmentation_mask, number_of_clues):
+        
+        lowest, highest = find_extreme_y_values(segmentation_mask)
+        if lowest is not None and highest is not None:
+                clue_step = int((highest - lowest) / number_of_clues+1)
+                clues = []
+                for i in range(number_of_clues):
+                        clues.append(highest - (i*clue_step))
+                clues.append(lowest+int(0.5*clue_step))
+                        
+                return clues
+        else:
+                return []
+
+def border_handler(id_map, image, edges, target_distances):
+        
+        lowest, _ = find_extreme_y_values(id_map)
+        borders = []
+        regions = []
+        for target in target_distances:
+                borders_regions = find_zone_border(id_map, image, edges, irl_target_mm=target, lowest_y = lowest)
+                borders.append(borders_regions[0])
+                regions.append(borders_regions[1])
+                
+        return borders, id_map, regions
+
+def segment(model_seg, image_size, filename, PATH_jpgs, dataset_type, model_type, item=None):
+        image_norm, _, image, mask, _ = load(filename, PATH_jpgs, image_size, dataset_type=dataset_type, item=item)
+        id_map = process(model_seg, image_norm, mask, model_type)
+        id_map = cv2.resize(id_map, [1920,1080], interpolation=cv2.INTER_NEAREST)
+        return id_map, image
+
+def detect(model_det, filename_img, PATH_jpgs):
+        
+        image = cv2.imread(os.path.join(PATH_jpgs, filename_img))
+        results = model_det.predict(image)
+
+        return results, model_det, image
+
+def manage_detections(results, model):
+        bbox = results[0].boxes.xywh.tolist()
+        cls = results[0].boxes.cls.tolist()
+        accepted_stationary = np.array([24,25,28,36])
+        accepted_moving = np.array([0,1,2,3,7,15,16,17,18,19])
+        boxes_moving = {}
+        boxes_stationary = {}
+        if len(bbox) > 0:
+                for xywh, clss in zip(bbox, cls):
+                        if clss in accepted_moving:
+                                if clss in boxes_moving.keys() and len(boxes_moving[clss]) > 0:
+                                        boxes_moving[clss].append(xywh)
+                                else:
+                                        boxes_moving[clss] = [xywh]
+                        if clss in accepted_stationary:
+                                if clss in boxes_stationary.keys() and len(boxes_stationary[clss]) > 0:
+                                        boxes_stationary[clss].append(xywh)
+                                else:
+                                        boxes_stationary[clss] = [xywh]
+
+        return boxes_moving, boxes_stationary
+
+def compute_detection_borders(borders, output_dims=[1080,1920]):
+        det_height = output_dims[0]-1
+        det_width = output_dims[1]-1
+        
+        for i,border in enumerate(borders):
+                border_l = np.array(border[0])
+                
+                if list(border_l):
+                        pass
+                else:
+                        border_l=np.array([[0,0],[0,0]])
+                
+                endpoints_l = [border_l[0],border_l[-1]]
+                
+                border_r = np.array(border[1])
+                if list(border_r):
+                        pass
+                else:
+                        border_r=np.array([[0,0],[0,0]])
+                        
+                endpoints_r = [border_r[0],border_r[-1]]
+                
+                if np.array_equal(np.array([[0,0],[0,0]]), endpoints_l):
+                        endpoints_l = [[0,endpoints_r[0][1]],[0,endpoints_r[1][1]]]
+                        
+                if np.array_equal(np.array([[0,0],[0,0]]), endpoints_r):
+                        endpoints_r = [[det_width,endpoints_l[0][1]],[det_width,endpoints_l[1][1]]]
+                
+                interpolated_top = bresenham_line(endpoints_l[1][0],endpoints_l[1][1],endpoints_r[1][0],endpoints_r[1][1])
+
+                zero_range = [0,1,2,3]
+                height_range = [det_height,det_height-1,det_height-2,det_height-3]
+                width_range = [det_width,det_width-1,det_width-2,det_width-3]
+
+                if (endpoints_l[0][0] in zero_range and endpoints_r[0][1] in height_range):
+                        y_values = np.arange(endpoints_l[0][1], det_height)
+                        x_values = np.full_like(y_values, 0)
+                        bottom1 = np.column_stack((x_values, y_values))
+                        
+                        x_values = np.arange(0, endpoints_r[0][0])
+                        y_values = np.full_like(x_values, det_height)
+                        bottom2 = np.column_stack((x_values, y_values))
+                        
+                        interpolated_bottom = np.vstack((bottom1, bottom2))
+                        
+                elif (endpoints_l[0][1] in height_range and endpoints_r[0][0] in width_range):
+                        y_values = np.arange(endpoints_r[0][1], det_height)
+                        x_values = np.full_like(y_values, det_width)
+                        bottom1 = np.column_stack((x_values, y_values))
+                        
+                        x_values = np.arange(endpoints_l[0][0], det_width)
+                        y_values = np.full_like(x_values, det_height)
+                        bottom2 = np.column_stack((x_values, y_values))
+                        
+                        interpolated_bottom = np.vstack((bottom1, bottom2))
+                        
+                elif endpoints_l[0][0] in zero_range and endpoints_r[0][0] in width_range:
+                        y_values = np.arange(endpoints_l[0][1], det_height)
+                        x_values = np.full_like(y_values, 0)
+                        bottom1 = np.column_stack((x_values, y_values))
+                        
+                        y_values = np.arange(endpoints_r[0][1], det_height)
+                        x_values = np.full_like(y_values, det_width)
+                        bottom2 = np.column_stack((x_values, y_values))
+                        
+                        bottom3_mid = bresenham_line(bottom1[-1][0],bottom1[-1][1],bottom2[-1][0],bottom2[-1][1])
+                        
+                        interpolated_bottom = np.vstack((bottom1, bottom2, bottom3_mid))
+
+                        
+                else:
+                        interpolated_bottom = bresenham_line(endpoints_l[0][0],endpoints_l[0][1],endpoints_r[0][0],endpoints_r[0][1])
+                
+                borders[i].append(interpolated_bottom)
+                borders[i].append(interpolated_top)
+                
+        return borders
+
+def get_bounding_box_points(cx, cy, w, h):
+        top_left = (cx - w / 2, cy - h / 2)
+        top_right = (cx + w / 2, cy - h / 2)
+        bottom_right = (cx + w / 2, cy + h / 2)
+        bottom_left = (cx - w / 2, cy + h / 2)
+        
+        corners = [top_left, top_right, bottom_right, bottom_left]
+        
+        def interpolate(point1, point2, fraction):
+                """Interpolate between two points at a given fraction of the distance."""
+                return (point1[0] + fraction * (point2[0] - point1[0]), 
+                        point1[1] + fraction * (point2[1] - point1[1]))
+
+        points = []
+        for i in range(4):
+                next_i = (i + 1) % 4
+                points.append(corners[i])
+                points.append(interpolate(corners[i], corners[next_i], 1 / 3))
+                points.append(interpolate(corners[i], corners[next_i], 2 / 3))
+
+        return points
+
+def classify_detections(boxes_moving, boxes_stationary, borders, img_dims, output_dims=[1080,1920]):
+        img_h, img_w, _ = img_dims
+        img_h_scaletofullHD = output_dims[1]/img_w
+        img_w_scaletofullHD = output_dims[0]/img_h
+        colors = ["yellow","orange","red","green","blue"]
+        
+        borders = compute_detection_borders(borders,output_dims)
+        
+        boxes_info = []
+        
+        if boxes_moving or boxes_stationary:
+                if boxes_moving:
+                        for item, coords in boxes_moving.items():
+                                for coord in coords:
+                                        x = coord[0]*img_w_scaletofullHD
+                                        y = coord[1]*img_h_scaletofullHD
+                                        w = coord[2]*img_w_scaletofullHD
+                                        h = coord[3]*img_h_scaletofullHD
+                                        
+                                        points_to_test = get_bounding_box_points(x, y, w, h)
+                                        
+                                        complete_border = []
+                                        criticality = -1
+                                        color = None
+                                        for i,border in enumerate(reversed(borders)):
+                                                border_nonempty = [np.array(arr) for arr in border if np.array(arr).size > 0]
+                                                complete_border = np.vstack((border_nonempty))
+                                                instance_border_path = mplPath.Path(np.array(complete_border))
+                                                
+                                                is_inside_borders = False
+                                                for point in points_to_test:
+                                                        is_inside = instance_border_path.contains_point(point)
+                                                        if is_inside:
+                                                                is_inside_borders = True
+                                                
+                                                if is_inside_borders:
+                                                        criticality = i
+                                                        color = colors[i]
+                                                        
+                                        if criticality == -1:
+                                                color = colors[3]
+                                                
+                                        boxes_info.append([item, criticality, color, [x, y], [w, h], 1])
+                                                
+                if boxes_stationary:
+                        for item, coords in boxes_stationary.items():
+                                for coord in coords:
+                                        x = coord[0]*img_w_scaletofullHD
+                                        y = coord[1]*img_h_scaletofullHD
+                                        w = coord[2]*img_w_scaletofullHD
+                                        h = coord[3]*img_h_scaletofullHD
+                                        
+                                        points_to_test = get_bounding_box_points(x, y, w, h)
+                                        
+                                        complete_border = []
+                                        criticality = -1
+                                        color = None
+                                        is_inside_borders = 0
+                                        for i,border in enumerate(reversed(borders), start=len(borders) - 1):
+                                                border_nonempty = [np.array(arr) for arr in border if np.array(arr).size > 0]
+                                                complete_border = np.vstack(border_nonempty)
+                                                instance_border_path = mplPath.Path(np.array(complete_border))
+                                                
+                                                is_inside_borders = False
+                                                for point in points_to_test:
+                                                        is_inside = instance_border_path.contains_point(point)
+                                                        if is_inside:
+                                                                is_inside_borders = True
+                                                
+                                                if is_inside_borders:
+                                                        criticality = i
+                                                        color = colors[4]
+                                                
+                                        if criticality == -1:
+                                                color = colors[3]
+                                                
+                                        boxes_info.append([item, criticality, color, [x, y], [w, h], 0])
+        
+                return boxes_info
+        
+        else:
+                print("No accepted detections in this image.")
+                return []
+
+def draw_classification(classification, id_map):
+        if classification:                
+                for box in classification:
+                        x,y = box[3]
+                        mark_value = 30
+                        
+                        x_start = int(max(x - 2, 0))
+                        x_end = int(min(x + 3, id_map.shape[1]))
+                        y_start = int(max(y - 2, 0))
+                        y_end = int(min(y + 3, id_map.shape[0]))
+                        
+                        id_map[y_start:y_end, x_start:x_end] = mark_value
+        else:
+                return
+
+def show_result(classification, id_map, names, borders, image, regions, file_index):
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        image = cv2.resize(image, (id_map.shape[1], id_map.shape[0]), interpolation = cv2.INTER_LINEAR)
+        fig = plt.figure(figsize=(16, 9), dpi=100)
+        plt.imshow(image, cmap='gray')
+        
+        if classification:
+                for box in classification:
+                        
+                        boxes = True
+                        cx,cy = box[3]
+                        name = names[box[0]]
+                        if boxes:
+                                w,h = box[4]
+                                x = cx - w / 2
+                                y = cy - h / 2
+                                rect = patches.Rectangle((x, y), w, h, linewidth=2, edgecolor=box[2], facecolor='none')
+                                
+                                ax = plt.gca()
+                                ax.add_patch(rect)
+                                plt.text(x, y-17, name, color='black', fontsize=10, ha='center', va='center', fontweight='bold', bbox=dict(facecolor=box[2], edgecolor='none', alpha=1))
+                        else:
+                                plt.imshow(id_map, cmap='gray')
+                                plt.text(cx, cy+10, name, color=box[2], fontsize=10, ha='center', va='center', fontweight='bold')
+
+        for region in regions:
+                for side in region:
+                        for line in side:
+                                line = np.array(line)
+                                plt.plot(line[:,1], line[:,0] ,'-', color='lightgrey', marker=None, linewidth=0.5)
+                                plt.ylim(0, 1080)
+                                plt.xlim(0, 1920)
+                                plt.gca().invert_yaxis()
+
+        colors = ['yellow','orange','red']
+        borders.reverse()
+        for i,border in enumerate(borders):
+                for side in border:
+                        side = np.array(side)
+                        if side.size > 0:
+                                plt.plot(side[:,0],side[:,1] ,'-', color=colors[i], marker=None, linewidth=0.6) #color=colors[i]
+                                plt.ylim(0, 1080)
+                                plt.xlim(0, 1920)
+                                plt.gca().invert_yaxis()
+                
+        plt.show()
+        #plt.tight_layout()
+        #plt.savefig(f'Grafika/Video_export/frames_estimated/frame_{file_index:04d}.jpg', format='jpg', bbox_inches='tight')
+        #plt.close()
+        print('Frame processed successfully.')
+
+def run(model_seg, model_det, image_size, filepath_img, PATH_jpgs, dataset_type, model_type, target_distances, file_index, vis, item=None, num_ys = 15):
+
+        segmentation_mask, image = segment(model_seg, image_size, filepath_img, PATH_jpgs, dataset_type, model_type, item)
+        print('File: {}'.format(filepath_img))
+        
+        # Border search
+        clues = get_clues(segmentation_mask, num_ys)
+        #edges = find_edges(segmentation_mask, clues, min_width=int(segmentation_mask.shape[1]*0.02))
+        edges = find_edges(segmentation_mask, clues, min_width=0)
+        #id_map_marked = mark_edges(segmentation_mask, edges)
+        
+        borders, id_map, regions = border_handler(segmentation_mask, image, edges, target_distances)
+        
+        # Detection
+        results, model, image = detect(model_det, filepath_img, PATH_jpgs)
+        boxes_moving, boxes_stationary = manage_detections(results, model)
+        
+        classification = classify_detections(boxes_moving, boxes_stationary, borders, image.shape, output_dims=segmentation_mask.shape)
+        
+        #draw_classification(classification, id_map)
+        show_result(classification, id_map, model.names, borders, image, regions, file_index)
+
+if __name__ == "__main__":
+
+        data_type = 'railsem19' #railsem19, pilsen or testdata
+        model_type = "segformer" #segformer or deeplab
+        vis = False
+        image_size = [1024,1024]
+        target_distances = [650,1000,2000] #[600,1000,2000] [4000,5500,6500] [2000,3000,4000]
+        num_ys = 10
+        
+        if data_type == 'pilsen':
+                file_index = 0
+                model_seg = load_model(PATH_model_seg)
+                model_det = load_yolo(PATH_model_det)
+                for item in enumerate(data_json["data"]):
+                        filepath_img = item[1][1]["path"]
+                        run(model_seg, model_det, image_size, filepath_img, PATH_base, data_type, model_type, target_distances, file_index, vis=vis, item=item, num_ys=num_ys)
+        elif data_type == 'railsem19':
+                file_index = 0
+                model_seg = load_model(PATH_model_seg)
+                model_det = load_yolo(PATH_model_det)
+                for filename_img in os.listdir(PATH_jpgs):
+                        #filename_img = "rs07650.jpg"
+                        run(model_seg, model_det, image_size, filename_img, PATH_jpgs, data_type, model_type, target_distances, file_index, vis=vis, item=None, num_ys=num_ys)
+                        file_index += 1
+        else:
+                file_index = 0
+                PATH_jpgs = 'Grafika/Video_export/frames'
+                model_seg = load_model(PATH_model_seg)
+                model_det = load_yolo(PATH_model_det)
+                for filename_img in os.listdir(PATH_jpgs):
+                        if os.path.exists(os.path.join('Grafika/Video_export/frames_estimated', filename_img)):
+                                file_index += 1
+                                continue
+                        else:
+                                run(model_seg, model_det, image_size  , filename_img, PATH_jpgs, data_type, model_type, target_distances, file_index, vis=vis, item=None, num_ys=num_ys)
+                                file_index += 1 

test_filtered_cls.py

@@ -0,0 +1,283 @@
+import numpy as np
+import pandas as pd
+import torch
+import cv2
+import os
+import torch.nn as nn
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+import torch.nn.functional as F
+from scripts.metrics_filtered_cls import compute_map_cls, compute_IoU, image_morpho
+from rs19_val.example_vis import rs19_label2bgr
+
+PATH_jpgs = 'assets/rs19val/jpgs/test'
+PATH_masks = 'assets/rs19val/uint8/test'
+PATH_model = 'assets/models_pretrained/segformer/SegFormer_B3_1024_finetuned.pth'
+
+def load(filename, PATH_jpgs, input_size=[224,224], dataset_type='rs19val', item = None):
+    transform_img = A.Compose([
+                    A.Resize(height=input_size[0], width=input_size[1], interpolation=cv2.INTER_NEAREST),
+                    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),
+                    ToTensorV2(p=1.0),
+                    ])
+    transform_mask = A.Compose([
+                    A.Resize(height=input_size[0], width=input_size[1], interpolation=cv2.INTER_NEAREST),
+                    ToTensorV2(p=1.0),
+                    ])  
+    
+    if dataset_type == 'pilsen':
+        mask_pth = item[1][1]["masks"]["ground_truth"]["path"]
+        mask_pth = os.path.join(PATH_jpgs, mask_pth)
+    elif dataset_type == 'railsem19':
+        mask_pth = os.path.join(PATH_masks, filename).replace('.jpg', '.png')
+    else:
+        mask_pth = "rs19_val/jpgs/placeholder_mask.png"
+        
+    image_in = cv2.imread(os.path.join(PATH_jpgs, filename))
+    mask = cv2.imread(mask_pth, cv2.IMREAD_GRAYSCALE)
+    
+    if dataset_type == 'testdata':
+        image_in = cv2.resize(image_in, (1920, 1080))
+    
+    image_tr = transform_img(image=image_in)['image']
+    image_tr = image_tr.unsqueeze(0)
+    image_vis = transform_mask(image=image_in)['image']
+    mask = transform_mask(image=mask)['image']
+    mask_id_map = np.array(mask.cpu().detach().numpy(), dtype=np.uint8)
+    
+    image_tr = image_tr.cpu()
+    
+    return image_tr, image_vis, image_in, mask, mask_id_map
+
+def load_model(path_model):
+    
+    model = torch.load(path_model, map_location=torch.device('cpu'))
+    model = model.cpu()
+    model.eval()
+    return model
+
+def remap_ignored_clss(id_map):
+    ignore_list = [0,1,2,6,8,9,15,16,19,20]
+    for cls in ignore_list:
+        id_map[id_map==cls] = 255
+
+    ignore_set = set(ignore_list)
+    cls_remaining = [num for num in range(0, 22) if num not in ignore_set]
+
+    # renumber the remaining classes 0-number of remaining classes
+    for idx, cls in enumerate(cls_remaining):
+        id_map[id_map==cls] = idx
+
+    id_map[id_map==255] = 12 # background
+    
+    return id_map
+
+def prepare_for_display(mask, image, id_map, rs19_label2bgr, image_size = [224,224]):
+    # Mask + prediction preparation
+    mask = mask + 1
+    mask[mask==256] = 0
+    mask = remap_ignored_clss(mask)
+    mask = (mask + 100).detach().numpy().squeeze().astype(np.uint8)
+    mask_rgb = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)
+
+    # Opacity channel addition to both mask and img
+    alpha_channel = np.full((mask.shape[0], mask.shape[1]), 255, dtype=np.uint8)
+    back_ids = mask==100
+    alpha_channel[back_ids] = 0
+    rgba_mask = cv2.merge((mask_rgb, alpha_channel))
+
+    image = np.array(image.cpu().detach().numpy(), dtype=np.uint8)
+    rgba_img = cv2.merge((image.transpose(1, 2, 0), alpha_channel))
+    
+    # Label colors + background
+    rgbs = list(rs19_label2bgr.values())
+    rgbs.append((255,255,255))
+
+    blend_sources = np.zeros((image_size[0], image_size[1], 3), dtype=np.uint8)
+    for class_id in range(21):
+        class_pixels = id_map == class_id
+        rgb_color = np.array(rgbs[class_id])
+
+        for i in range(3):
+            blend_sources[:,:,i] = blend_sources[:,:,i] + (rgb_color[i] * class_pixels).astype(np.uint8)
+
+    # Opacity channel for the rgb class mask and merge with input mask
+    alpha_channel_blend = np.full((blend_sources.shape[0], blend_sources.shape[1]), 150, dtype=np.uint8)
+    rgba_blend = cv2.merge((blend_sources , alpha_channel_blend))
+    blend_sources = (rgba_blend * 0.1 + rgba_img * 0.9).astype(np.uint8)
+    
+    return(rgba_mask, rgba_blend, blend_sources)
+    
+def visualize(rgba_blend, rgba_mask):
+    # CV2 VIZUALISATION
+    image1 = rgba_blend
+    image2 = rgba_mask
+
+    initial_opacity1 = 0.05
+    initial_opacity2 = 0.95
+    # Load two smaller images
+    small_image1 = cv2.resize(image1, (300, 300), interpolation=cv2.INTER_NEAREST)
+    small_image2 = cv2.resize(image2, (300, 300), interpolation=cv2.INTER_NEAREST)
+        
+    # Create a blank canvas for the combined visualization
+    combined_image = np.zeros((600, 900, 4), dtype=np.uint8)  # Adjust the size as needed
+
+    # Main loop for adjusting opacity and displaying the images
+    cv2.namedWindow('{} | mAP:{:.3f} | MmAP:{:.3f} '.format(filename, map, Mmap), cv2.WINDOW_NORMAL)
+    cv2.resizeWindow('{} | mAP:{:.3f} | MmAP:{:.3f} '.format(filename, map, Mmap), 900, 600)  # Adjust the size as needed
+
+    while True:
+        
+        overlay_image = image1.copy()
+        overlay_image[:, :, 3] = (image1[:, :, 3] * initial_opacity1).astype(np.uint8)
+
+        alpha = (image2[:, :, 3] * initial_opacity2).astype(float)
+        beta = 1.0 - alpha / 255.0
+
+        blended_image = np.empty_like(overlay_image)
+        blended_image[:, :, :3] = (overlay_image[:, :, :3] * alpha[:, :, np.newaxis] + image2[:, :, :3] * beta[:, :, np.newaxis]).astype(np.uint8)
+        blended_image[:, :, 3] = (overlay_image[:, :, 3] * alpha + image2[:, :, 3] * beta).astype(np.uint8)
+
+        blended_image = (image1 * initial_opacity1 + image2 * initial_opacity2).astype(np.uint8)
+
+        blended_image_resized = cv2.resize(blended_image, (600, 600))  # Adjust the size as needed
+        combined_image[:, :600, :] = blended_image_resized
+
+        # Copy the smaller images to the right portion of the canvas
+        combined_image[0:300, 600:900, :] = small_image1[:, :, :]
+        combined_image[300:600, 600:900, :] = small_image2[:, :, :]
+
+        cv2.imshow('{} | mAP:{:.3f} | MmAP:{:.3f} '.format(filename, map, Mmap), combined_image)
+
+        key = cv2.waitKey(1) & 0xFF
+        if key == ord('q'):
+            break
+        elif key == ord('a'):
+            initial_opacity1 += 0.1
+            initial_opacity1 = min(initial_opacity1, 1.0)
+        elif key == ord('s'):
+            initial_opacity1 -= 0.1
+            initial_opacity1 = max(initial_opacity1, 0.0)
+        elif key == ord('z'):
+            initial_opacity2 += 0.1
+            initial_opacity2 = min(initial_opacity2, 1.0)
+        elif key == ord('x'):
+            initial_opacity2 -= 0.1
+            initial_opacity2 = max(initial_opacity2, 0.0)
+
+    cv2.destroyAllWindows()
+
+def stats_mean_and_reorder(classes_ap,classes_Map,classes_stats,classes_Mstats):
+    for cls, value in classes_ap.items():
+        classes_ap[cls] = np.divide(value[0], value[1])
+    classes_ap['all']= np.mean(np.array(list(classes_ap.values())), axis=0)
+
+    for cls, value in classes_Map.items():
+        classes_Map[cls] = np.divide(value[0], value[1])
+    classes_Map['all']= np.mean(np.array(list(classes_Map.values())), axis=0)
+
+    for cls, value in classes_stats.items():
+        classes_stats[cls] = np.divide(value[0], value[1])
+    classes_stats['all']= np.mean(np.array(list(classes_stats.values()))[:, :4], axis=0)
+
+    for cls, value in classes_Mstats.items():
+        classes_Mstats[cls] = np.divide(value[0], value[1])
+    classes_Mstats['all']= np.mean(np.array(list(classes_Mstats.values()))[:, :4], axis=0)
+
+    for cls, value in classes_Mstats.items():
+        classes_stats[cls] = np.insert(classes_stats[cls], 1, value[0])
+        classes_stats[cls] = np.insert(classes_stats[cls], 3, value[1])
+        classes_stats[cls] = np.insert(classes_stats[cls], 5, value[2])
+        classes_stats[cls] = np.insert(classes_stats[cls], 7, value[3])
+
+    return classes_ap,classes_Map,classes_stats,classes_Mstats
+
+def process(model, input_img, mask, model_type):
+    if model_type == "segformer":
+        outputs = model(input_img) # segformer
+    elif model_type == "deeplab":
+        outputs = model(input_img)['out'] # deeplab resnet
+    
+    logits = outputs.logits
+    upsampled_logits = nn.functional.interpolate(
+        logits,
+        size=mask.shape[-2:],
+        mode="bilinear",
+        align_corners=False
+    )
+    
+    output  = upsampled_logits.float()
+        
+    confidence_scores = F.softmax(output, dim=1).cpu().detach().numpy().squeeze()
+    id_map = np.argmax(confidence_scores, axis=0).astype(np.uint8)
+    id_map = image_morpho(id_map)
+    
+    return id_map
+
+if __name__ == "__main__":
+    mAPs,MmAPs,IoUs,MIoUs,accs,Maccs,precs,Mprecs,recs,Mrecs= list(),list(),list(),list(),list(),list(),list(),list(),list(),list()
+    classes_ap,classes_Map,classes_stats,classes_Mstats = {},{},{},{}
+    images_computed = 0
+    
+    for filename in os.listdir(PATH_jpgs):
+        images_computed += 1
+        
+        vis = False
+        to_break = False
+        image_size = [1024,1024]
+        
+        if to_break:
+            if images_computed > 50:
+                break
+        
+        model_type = "segformer" #"deeplab"
+        dataset_type = 'rs19val'
+        image_norm, image, _, mask, id_map_gt = load(filename, PATH_jpgs, image_size, dataset_type)
+        model = load_model(PATH_model)
+        # INFERENCE + SOFTMAX
+        id_map = process(model, image_norm, mask, model_type)
+        
+        # mAP
+        id_map_gt = remap_ignored_clss(id_map_gt)
+        map,classes_ap  = compute_map_cls(id_map_gt, id_map, classes_ap)
+        Mmap,classes_Map = compute_map_cls(id_map_gt, id_map, classes_Map, major = True)
+        IoU,acc,prec,rec,classes_stats = compute_IoU(id_map_gt, id_map, classes_stats)
+        MIoU,Macc,Mprec,Mrec,classes_Mstats = compute_IoU(id_map_gt, id_map, classes_Mstats, major=True)
+        
+        print('{} | mAP:{:.3f}/{:.3f} | IoU:{:.3f}/{:.3f} | prec:{:.3f}/{:.3f} | rec:{:.3f}/{:.3f} | acc:{:.3f}/{:.3f}'.format(filename,map,Mmap,IoU,MIoU,prec,Mprec,rec,Mrec,acc,Macc))
+        mAPs.append(map)
+        MmAPs.append(Mmap)
+        IoUs.append(IoU)
+        MIoUs.append(MIoU)
+        accs.append(acc)
+        Maccs.append(Macc)
+        precs.append(prec)
+        Mprecs.append(Mprec)
+        recs.append(rec)
+        Mrecs.append(Mrec)
+        
+        if vis:
+            rgba_mask, rgba_blend, blend_sources = prepare_for_display(mask, image, id_map, rs19_label2bgr, image_size)        
+            visualize(rgba_blend, rgba_mask)
+
+    mAPs_avg, MmAPs_avg = np.nanmean(mAPs), np.nanmean(MmAPs)
+    IoUs_avg, MIoUs_avg = np.nanmean(IoUs), np.nanmean(MIoUs)
+    accs_avg, Maccs_avg = np.nanmean(accs), np.nanmean(Maccs)
+    precs_avg, Mprecs_avg = np.nanmean(precs), np.nanmean(Mprecs)
+    recs_avg, Mrecs_avg = np.nanmean(recs), np.nanmean(Mrecs)
+
+    print('All         | mAP:{:.3f}/{:.3f} | IoU:{:.3f}/{:.3f} | prec:{:.3f}/{:.3f} | rec:{:.3f}/{:.3f} | acc:{:.3f}/{:.3f}'.format(mAPs_avg,MmAPs_avg,IoUs_avg,MIoUs_avg,precs_avg,Mprecs_avg,recs_avg,Mrecs_avg,accs_avg,Maccs_avg))
+    print('mAP: {:.3f}-{:.3f} | MmAP: {:.3f}-{:.3f} | IoU: {:.3f}-{:.3f} | MIoU: {:.3f}-{:.3f}'.format(np.nanmin(mAPs), np.nanmax(mAPs), np.nanmin(MmAPs), np.nanmax(MmAPs),np.nanmin(IoUs), np.nanmax(IoUs), np.nanmin(MIoUs), np.nanmax(MIoUs)))
+
+    classes_ap,classes_Map,classes_stats,classes_Mstats = stats_mean_and_reorder(classes_ap,classes_Map,classes_stats,classes_Mstats)
+
+    df_ap = pd.DataFrame(list(classes_ap.items()), columns=['Class', 'mAP'])
+    df_Map = pd.DataFrame(list(classes_Map.items()), columns=['Class', 'MmAP'])
+
+    classes_stats_flat = [(key, *value) for key, value in classes_stats.items()]
+    df_stats = pd.DataFrame(classes_stats_flat, columns=['Class','IoU','MIoU', 'acc','Macc', 'precision','Mprecision','recall','Mrecall'])
+
+    df_merged = pd.merge(df_ap, df_Map, on='Class', how='outer')
+    df_merged = pd.merge(df_merged, df_stats, on='Class', how='outer')
+
+    print(df_merged)