Source code for mactrack.locate.defuse

import cv2
import numpy as np
import os


[docs] def trace_lines_between_contours(images, distance_threshold=50): """ Trace lines between contours in a list of images based on a distance threshold. The contours are obtained from the binary images by the ``findContours`` function from the OpenCV library. Parameters: images (list): List of images (numpy arrays) to process. distance_threshold (int): Maximum distance between contours to draw a line. Returns: traced_lines_image (numpy array): Image with traced lines between contours. """ traced_lines_image = np.zeros_like(images[0]) all_contours = [] for image in images: contours, _ = cv2.findContours( image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE ) all_contours.extend(contours) for i in range(len(all_contours)): for j in range(i + 1, len(all_contours)): cnt1 = all_contours[i] cnt2 = all_contours[j] for point1 in cnt1: for point2 in cnt2: dist = np.linalg.norm(point1 - point2) if dist < distance_threshold: pt1 = tuple(point1[0]) pt2 = tuple(point2[0]) cv2.line(traced_lines_image, pt1, pt2, (255), 1) return traced_lines_image
[docs] def paint_black_area_from_mask(image_c, mask): """ Paints the black area of the image based on the mask. Parameters: image_c (numpy array): The original image (BGR or grayscale). mask (numpy array): The mask to use for painting the black area. Returns: image_c_black (numpy array): The image with the black area painted. """ if len(image_c.shape) == 2: image_c_gray = image_c elif len(image_c.shape) == 3: image_c_gray = cv2.cvtColor(image_c, cv2.COLOR_BGR2GRAY) else: raise ValueError( "Unsupported image format. Expected grayscale (1 channel) or BGR (3 channels)." ) inverted_mask = cv2.bitwise_not(mask) black_area = cv2.bitwise_and(image_c_gray, image_c_gray, mask=inverted_mask) image_c_black = cv2.merge((black_area, black_area, black_area)) return image_c_black
[docs] def extract_and_save_objects( image_c, image, heatmap, object, image_storage, min_object_size=100 ): """ Extracts and saves objects from the image based on contours. Parameters: image_c (numpy array): The original image (BGR or grayscale). image (numpy array): The image with contours. heatmap (int): The heatmap number. object (int): The object number. image_storage (ImageStorage): The image storage object. min_object_size (int): Minimum size of the object to be extracted. Default is 100. Returns: image_storage (ImageStorage): The updated image storage object. """ image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) contours, _ = cv2.findContours( image_gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE ) replaced = False n = len(image_storage.get_list(f"heatmap_test_{heatmap}")) j = 0 for i, contour in enumerate(contours): if cv2.contourArea(contour) < min_object_size: continue mask = np.zeros_like(image_gray) object_image = cv2.drawContours( mask, [contour], -1, (255, 255, 255), thickness=cv2.FILLED ) object_image = cv2.bitwise_and(image_c, image_c, mask=mask) if not replaced: image_storage.replace_image( f"heatmap_test_{heatmap}", f"object_{object}.png", object_image ) replaced = True else: image_storage.add_image( f"heatmap_test_{heatmap}", f"object_{n+j}.png", object_image ) j = j + 1 return image_storage
[docs] def process_images(list, image_c, heatmap, object, image_storage, max_line_length=50): """ Process images by tracing lines between contours and painting the black area. Parameters: list (list): List of images (numpy arrays) to process. image_c (numpy array): The original image (BGR or grayscale). heatmap (int): The heatmap number. object (int): The object number. image_storage (ImageStorage): The image storage object. max_line_length (int): Maximum length of the line to be drawn. Default is 50. Returns: image_storage (ImageStorage): The updated image storage object. """ n = len(list) print(n) traced_lines_image = trace_lines_between_contours(list) result_image = np.zeros_like(traced_lines_image) result_image = cv2.bitwise_or(result_image, traced_lines_image) for i in range(n): contours, _ = cv2.findContours( list[i], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE ) cv2.drawContours(result_image, contours, -1, (0, 0, 0), thickness=cv2.FILLED) image = paint_black_area_from_mask(image_c, result_image) extract_and_save_objects(image_c, image, heatmap, object, image_storage) return image_storage
[docs] def calculate_iou(image1, image2): """ Calculate the Intersection over Union (IoU) between two binary images. .. math:: IoU = \\frac{A \\cap B}{A \\cup B} where :math:`A` and :math:`B` are the two binary images. Parameters: image1 (numpy array): First binary image. image2 (numpy array): Second binary image. Returns: float: IoU value between 0 and 1. """ intersection = np.logical_and(image1, image2).sum() union = np.logical_or(image1, image2).sum() return intersection / union if union != 0 else 0
[docs] def save_segmentation_images(segmentation_instance, output_folder): """ Save the segmented images to the specified output folder. Parameters: segmentation_instance (Segmentation): The segmentation instance containing the images. output_folder (str): The folder where the images will be saved. """ if not os.path.exists(output_folder): os.makedirs(output_folder) for dir_name, objects_list in segmentation_instance.images: heatmap_folder = os.path.join(output_folder, dir_name) os.makedirs(heatmap_folder, exist_ok=True) for file_name, sparse_image in objects_list: dense_image = sparse_image.toarray().astype(np.uint8) image_path = os.path.join(heatmap_folder, file_name) cv2.imwrite(image_path, dense_image)
[docs] def defuse(n, image_storage): """ Splits the images by processing them and saving the results. Parameters: n (int): Number of images to process. image_storage (ImageStorage): The image storage object. Returns: image_storage (ImageStorage): The updated image storage object. """ c = 0 for i in range(1, n): print(f"heatmap_test_{i}") current = len(image_storage.get_list(f"heatmap_test_{i}")) for j in range(current): image1 = image_storage.get_image(f"heatmap_test_{i}", f"object_{j}.png") if image1 is not None: image1 = image1.toarray() _, image1 = cv2.threshold(image1, 127, 255, cv2.THRESH_BINARY) else: print(i, j) matches = [] previous = len(image_storage.get_list(f"heatmap_test_{i-1}")) for k in range(previous): image2 = image_storage.get_image( f"heatmap_test_{i-1}", f"object_{k}.png" ) if image2 is not None: image2 = image2.toarray() _, image2 = cv2.threshold(image2, 127, 255, cv2.THRESH_BINARY) else: print(i - 1, k) iou = calculate_iou(image1, image2) if iou > 0: matches.append(image2) c += 1 if c > 1: print(i, j) image_storage = process_images(matches, image1, i, j, image_storage) c = 0 return image_storage
[docs] def invdefuse(n, image_storage): """ Inverse splits the images by processing them and saving the results. Parameters: n (int): Number of images to process. image_storage (ImageStorage): The image storage object. Returns: image_storage (ImageStorage): The updated image storage object. """ caller_dir = os.path.dirname(os.path.abspath(__file__)) parent = os.path.dirname(caller_dir) root = os.path.dirname(parent) output_path = os.path.join(root, r"output") c = 0 for i in range(1, n): print(f"heatmap_test_{n-i-1}") current = len(image_storage.get_list(f"heatmap_test_{n-i-1}")) for j in range(current): image1 = image_storage.get_image(f"heatmap_test_{n-i-1}", f"object_{j}.png") if image1 is not None: image1 = image1.toarray() _, image1 = cv2.threshold(image1, 127, 255, cv2.THRESH_BINARY) else: print(n - i - 1, j) matches = [] previous = len(image_storage.get_list(f"heatmap_test_{n-i}")) for k in range(previous): image2 = image_storage.get_image( f"heatmap_test_{n-i}", f"object_{k}.png" ) if image2 is not None: image2 = image2.toarray() _, image2 = cv2.threshold(image2, 127, 255, cv2.THRESH_BINARY) else: print(n - i, k) iou = calculate_iou(image1, image2) if iou > 0: matches.append(image2) c += 1 if c > 1: print(n - i - 1, j) image_storage = process_images( matches, image1, n - i - 1, j, image_storage ) c = 0 save_segmentation_images(image_storage, os.path.join(output_path, "list_def")) return image_storage