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