File size: 4,441 Bytes
15acbf0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import random
import math
import numpy as np
from PIL import Image
from skimage.draw import line
from skimage import morphology
import cv2
def line_crosses_cracks(start, end, img):
rr, cc = line(start[0], start[1], end[0], end[1])
# Exclude the starting point from the line coordinates
if len(rr) > 1 and len(cc) > 1:
return np.any(img[rr[1:], cc[1:]] == 255)
return False
def random_walk(img_array, k=8, m=0.1, min_steps=50, max_steps=200, length=2, degree_range=30, seed=None):
if seed is not None:
random.seed(seed)
np.random.seed(seed)
img_array = cv2.ximgproc.thinning(img_array)
rows, cols = img_array.shape
# Find all white pixels (existing cracks)
white_pixels = np.column_stack(np.where(img_array == 255))
original_crack_count = len(white_pixels) # Count of original crack pixels
# Select k random starting points from the white pixels
if white_pixels.size == 0:
raise ValueError("No initial crack pixels found in the image.")
if k > len(white_pixels):
raise ValueError("k is greater than the number of existing crack pixels.")
initial_points = white_pixels[random.sample(range(len(white_pixels)), k)]
# Initialize step count for each initial point with a random value between min_steps and max_steps
step_counts = {i: random.randint(min_steps, max_steps) for i in range(k)}
# Initialize main direction for each initial point (0 to 360 degrees)
main_angles = {i: random.uniform(0, 360) for i in range(k)}
grown_crack_count = 0 # Count of newly grown crack pixels
# Start the random walk for each initial point
for idx, point in enumerate(initial_points):
current_pos = tuple(point)
current_steps = 0
while current_steps < step_counts[idx]:
# Check the crack ratio
current_ratio = np.sum(img_array == 255) / (rows * cols)
if current_ratio >= m:
return img_array, {'original_crack_count': original_crack_count, 'grown_crack_count': grown_crack_count}
# Generate a random direction within the fan-shaped area around the main angle
main_angle = main_angles[idx]
angle = math.radians(main_angle + random.uniform(-degree_range, degree_range))
# Determine the next position with the specified length
delta_row = length * math.sin(angle)
delta_col = length * math.cos(angle)
next_pos = (int(current_pos[0] + delta_row), int(current_pos[1] + delta_col))
# Check if the line from the current to the next position crosses existing cracks
if 0 <= next_pos[0] < rows and 0 <= next_pos[1] < cols and not line_crosses_cracks(current_pos, next_pos, img_array):
# Draw a line from the current position to the next position
rr, cc = line(current_pos[0], current_pos[1], next_pos[0], next_pos[1])
img_array[rr, cc] = 255 # Set the pixels along the line to white
grown_crack_count += len(rr) # Update the count of grown crack pixels
current_pos = next_pos
current_steps += 1
else:
# If the line crosses existing cracks or the next position is outside the boundaries, stop the walk for this point
break
return img_array, {'original_crack_count': original_crack_count, 'grown_crack_count': grown_crack_count}
# The rest of the test code remains the same.
# You can use this function in your test code to generate the image and get the counts.
# test code
if __name__ == "__main__":
# Updated parameters
k = 8 # Number of initial white pixels to start the random walk
m = 0.1 # Maximum ratio of crack pixels
min_steps = 50
max_steps = 200
img_path = '/data/leiqin/diffusion/huggingface_diffusers/crack_label_creator/random_walk/thindata_256/2.png'
img = Image.open(img_path)
img_array = np.array(img)
length = 2
# Perform the modified random walk
result_img_array_mod, pixels_dict = random_walk(img_array.copy(), k, m, min_steps, max_steps, length)
# Convert the result to an image
result_img_mod = Image.fromarray(result_img_array_mod.astype('uint8'))
# Save the resulting image
result_img_path_mod = 'resutls.png'
result_img_mod.save(result_img_path_mod)
print(pixels_dict)
|