Skip to content

Commit

Permalink
creates overlay for both crackmask and filteredCrackMasks
Browse files Browse the repository at this point in the history
  • Loading branch information
@cnpcshangbo
cnpcshangbo committed Sep 11, 2024
1 parent 3a979ac commit 9197073
Show file tree
Hide file tree
Showing 5 changed files with 169 additions and 132 deletions.
14 changes: 12 additions & 2 deletions DirectoryImageMaskProcessor_2curve.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,18 @@ def process_directory(self):
for idx, red_mask_filename in enumerate(image_files, start=1):
print(f"Processing image {idx}/{total_images}: {red_mask_filename}")
# Construct the raw image path using the directory and filename
raw_image_path = os.path.join(self.raw_directory_path, red_mask_filename.replace(".png", ".JPG")) # Adjust the extension if necessary

raw_image_base = os.path.splitext(red_mask_filename)[0]
raw_image_path = None
for ext in ['.png', '.PNG', '.jpg', '.JPG', '.jpeg', '.JPEG']:
potential_path = os.path.join(self.raw_directory_path, raw_image_base + ext)
if os.path.exists(potential_path):
raw_image_path = potential_path
break

if raw_image_path is None:
print(f"Warning: No matching raw image found for {red_mask_filename}")
continue

self.process_image(raw_image_path, red_mask_filename)

print("All images processed successfully.")
Expand Down
171 changes: 87 additions & 84 deletions crackoverlay_transparent.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,93 +10,96 @@
config = configparser.ConfigParser()
config.read("config.ini") # Replace with the actual path to your INI file

# Paths for the mask images, raw images, and output images
# Read from the INI file



mask_dir = config["CrackSegmentation"]["mask_directory"].replace(
"crackmask", "filteredCrackMasks"
)
# Define paths for both types of masks and their corresponding output directories
mask_dirs = {
"crackmask": config["CrackSegmentation"]["mask_directory"],
"filteredCrackMasks": config["CrackSegmentation"]["mask_directory"].replace("crackmask", "filteredCrackMasks")
}
raw_dir = config["Settings"]["image_path"]
output_dir = config["CrackOverlay"]["overlay_directory"].replace(
"crackoverlay", "filteredCrackOverlays"
)

# Compatible with the situation when concrete filter is not used.
if not os.path.exists(mask_dir):
mask_dir = config["CrackSegmentation"]["mask_directory"]
output_dir = config["CrackOverlay"]["overlay_directory"]

# Ensure the output directory exists
if not os.path.exists(output_dir):
os.makedirs(output_dir)

# Iterate over the files in the mask directory
for mask_name in os.listdir(mask_dir):
# Extract the image name from the filename
raw_name = mask_name.split(".")[0] + ".jpg"
if not os.path.isfile(os.path.join(raw_dir, raw_name)):
raw_name = mask_name.split(".")[0] + ".JPG"

# Load the mask and raw images
mask = cv2.imread(os.path.join(mask_dir, mask_name), cv2.IMREAD_GRAYSCALE)
raw = cv2.imread(os.path.join(raw_dir, raw_name))

if mask is None or raw is None:
print(f"Error reading {mask_name} or {raw_name}. Skipping...")
output_dirs = {
"crackmask": config["CrackOverlay"]["overlay_directory"],
"filteredCrackMasks": config["CrackOverlay"]["overlay_directory"].replace("crackoverlay", "filteredCrackOverlays")
}

# Process each type of mask
for mask_type, mask_dir in mask_dirs.items():
if not os.path.exists(mask_dir):
print(f"Skipping {mask_type} processing: {mask_dir} does not exist.")
continue
# print(f"np.unique(binary_mask)= {np.unique(mask)}")
# exit()
# Convert the mask to a binary image
# binary_mask = cv2.threshold(mask, 38, 255, cv2.THRESH_BINARY)[1] # pixel = 38 means crack
# Threshold for crack pixels (value to be set to 255)
crack_threshold = 38

# Values to set to 0
zero_values = [0, 75]

# Create a new binary mask with desired values
binary_mask = np.zeros_like(mask) # Create a zero-filled mask of the same shape

# Set crack_threshold value to 255
binary_mask[mask == crack_threshold] = 255

# Set zero_values to 0
for value in zero_values:
binary_mask[mask == value] = 0
# Get the dimensions of the mask and raw images
mask_height, mask_width = binary_mask.shape[:2]
raw_height, raw_width = raw.shape[:2]

# Resize the mask to match the dimensions of the raw image
binary_mask = cv2.resize(binary_mask, (raw_width, raw_height))

# Create a weight map for the mask
# weight_map = cv2.addWeighted(binary_mask, 0.5, binary_mask, 0.5, 0)

# ... (Your code for loading images, etc.)

# Create a base image filled with black (no color initially)
weight_map = np.zeros((raw_height, raw_width, 3), dtype=np.uint8)

# Find the locations in the binary mask where there are cracks
crack_locations = np.where(binary_mask != 0)

# Set the crack areas to green in the weight map
weight_map[crack_locations] = (0, 120, 0) # BGR for a green color

# Now overlay the weight map directly onto the original image:
overlay = cv2.addWeighted(raw, 1, weight_map, 0.5, 0)

# Save the intermediate results
# cv2.imwrite(os.path.join(output_dir, f"{raw_name}_mask.jpg"), mask)
# cv2.imwrite(os.path.join(output_dir, f"{raw_name}_binary_mask.jpg"), binary_mask)
# cv2.imwrite(os.path.join(output_dir, f"{raw_name}_weight_map.jpg"), weight_map)
output_dir = output_dirs[mask_type]

# Ensure the output directory exists
if not os.path.exists(output_dir):
os.makedirs(output_dir)

# Save the overlaid image to the output directory
output_name = os.path.join(output_dir, raw_name)
cv2.imwrite(output_name, overlay)
print(f"Saved {output_name}")
print(f"Processing {mask_type}...")

# Iterate over the files in the mask directory
for mask_name in os.listdir(mask_dir):
# Extract the image name from the filename
raw_name = mask_name.split(".")[0] + ".jpg"
if not os.path.isfile(os.path.join(raw_dir, raw_name)):
raw_name = mask_name.split(".")[0] + ".JPG"

# Load the mask and raw images
mask = cv2.imread(os.path.join(mask_dir, mask_name), cv2.IMREAD_GRAYSCALE)
raw = cv2.imread(os.path.join(raw_dir, raw_name))

if mask is None or raw is None:
print(f"Error reading {mask_name} or {raw_name}. Skipping...")
continue
# print(f"np.unique(binary_mask)= {np.unique(mask)}")
# exit()
# Convert the mask to a binary image
# binary_mask = cv2.threshold(mask, 38, 255, cv2.THRESH_BINARY)[1] # pixel = 38 means crack
# Threshold for crack pixels (value to be set to 255)
crack_threshold = 38

# Values to set to 0
zero_values = [0, 75]

# Create a new binary mask with desired values
binary_mask = np.zeros_like(mask) # Create a zero-filled mask of the same shape

# Set crack_threshold value to 255
binary_mask[mask == crack_threshold] = 255

# Set zero_values to 0
for value in zero_values:
binary_mask[mask == value] = 0
# Get the dimensions of the mask and raw images
mask_height, mask_width = binary_mask.shape[:2]
raw_height, raw_width = raw.shape[:2]

# Resize the mask to match the dimensions of the raw image
binary_mask = cv2.resize(binary_mask, (raw_width, raw_height))

# Create a weight map for the mask
# weight_map = cv2.addWeighted(binary_mask, 0.5, binary_mask, 0.5, 0)

# ... (Your code for loading images, etc.)

# Create a base image filled with black (no color initially)
weight_map = np.zeros((raw_height, raw_width, 3), dtype=np.uint8)

# Find the locations in the binary mask where there are cracks
crack_locations = np.where(binary_mask != 0)

# Set the crack areas to green in the weight map
weight_map[crack_locations] = (0, 120, 0) # BGR for a green color

# Now overlay the weight map directly onto the original image:
overlay = cv2.addWeighted(raw, 1, weight_map, 0.5, 0)

# Save the intermediate results
# cv2.imwrite(os.path.join(output_dir, f"{raw_name}_mask.jpg"), mask)
# cv2.imwrite(os.path.join(output_dir, f"{raw_name}_binary_mask.jpg"), binary_mask)
# cv2.imwrite(os.path.join(output_dir, f"{raw_name}_weight_map.jpg"), weight_map)

# Save the overlaid image to the output directory
output_name = os.path.join(output_dir, raw_name)
cv2.imwrite(output_name, overlay)
print(f"Saved {output_name}")

print("Processing complete!")
20 changes: 13 additions & 7 deletions main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -94,26 +94,32 @@ def main():
# Run crack related processing
log_message("Running crack segmentation...")
call_in_conda_env("python cracksegmentation.py")
# Produces: run_timestamp/mask/crack_mask

# log_message("Running concrete mask...")
# call_in_conda_env("python concretemask.py")
# Produces: run_timestamp/mask/concrete_mask

if config['concrete_post_filter']:
log_message("Running concrete post filter...")
call_in_conda_env("python concretePostFilter.py")
# Updates: run_timestamp/mask/crack_mask

log_message("Running crack23directions...")
log_message("Converting crack masks to 3 categories according to directions...")
call_in_conda_env("python crack23directions.py")
# Produces: run_timestamp/mask/crack_mask_3directions

log_message("Running crack2curve...")
call_in_conda_env("python crack2curve.py")
# Produces: run_timestamp/mask/crack_curve

exit()

log_message("Converting crack masks to 3 categories according to directions...")
call_in_conda_env("python crack23directions.py")
# Export nnfilteredCrackOverlay
log_message("Export nnfilteredCrackOverlay...")
call_in_conda_env("python export_nnfilteredCrackOverlay.py")

log_message("Running crack overlay...")
call_in_conda_env("python crackoverlay_transparent.py")
# Produces: run_timestamp/overlay/crack_overlay and filteredCrackOverlay and nnfilteredCrackOverlay?

# log_message("Copying geolocation info to crack overlay...")
# call_in_conda_env("python copy_geolocation_crack.py")
Expand Down Expand Up @@ -147,8 +153,8 @@ def main():
call_in_conda_env("python stainoverlay_transparent.py")
exit()

log_message("Copying geolocation info to stain overlay...")
call_in_conda_env("python3 copy_geolocation_stain.py")
# log_message("Copying geolocation info to stain overlay...")
# call_in_conda_env("python3 copy_geolocation_stain.py")

# log_message("Convert overlay images to pointcloud. ")
# call_in_conda_env("python3 overlay2pointcloud.py --damage_type stain")
Expand Down
91 changes: 56 additions & 35 deletions stainoverlay_transparent.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,50 +10,71 @@

# Paths for the mask images, raw images, and output images
# Read from the INI file
mask_dir = config["StainSegmentation"]["mask_directory"].replace("stainmask", "filteredStainMasks")
stain_mask_dir = config["StainSegmentation"]["mask_directory"]
filtered_mask_dir = stain_mask_dir.replace("stainmask", "filteredStainMasks")
raw_dir = config["Settings"]["image_path"]
output_dir = config["StainOverlay"]["overlay_directory"].replace("stainoverlay", "filteredStainOverlays")
stain_output_dir = config["StainOverlay"]["overlay_directory"]
filtered_output_dir = stain_output_dir.replace("stainoverlay", "filteredStainOverlays")

# Ensure the output directory exists
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# Function to process masks
def process_masks(mask_dir, output_dir):
# Ensure the directories exist
for directory in [mask_dir, output_dir]:
if not os.path.exists(directory):
print(f"Creating directory: {directory}")
os.makedirs(directory)

# Iterate over the files in the mask directory
for mask_name in os.listdir(mask_dir):
# Extract the image name from the filename
raw_name = mask_name.split(".")[0] + ".jpg"
if not os.path.isfile(os.path.join(raw_dir, raw_name)):
raw_name = mask_name.split(".")[0] + ".JPG"
# Check if the mask directory is empty
if not os.listdir(mask_dir):
print(f"Warning: Mask directory '{mask_dir}' is empty. No files to process.")
return

# Load the mask and raw images
mask = cv2.imread(os.path.join(mask_dir, mask_name), cv2.IMREAD_GRAYSCALE)
raw = cv2.imread(os.path.join(raw_dir, raw_name))
# Iterate over the files in the mask directory
for mask_name in os.listdir(mask_dir):
# Extract the image name from the filename
raw_name = mask_name.split(".")[0] + ".jpg"
if not os.path.isfile(os.path.join(raw_dir, raw_name)):
raw_name = mask_name.split(".")[0] + ".JPG"

if mask is None or raw is None:
print(f"Error reading {mask_name} or {raw_name}. Skipping...")
continue
# Load the mask and raw images
mask = cv2.imread(os.path.join(mask_dir, mask_name), cv2.IMREAD_GRAYSCALE)
raw = cv2.imread(os.path.join(raw_dir, raw_name))

print(raw.shape)
print(mask.shape)
if mask is None or raw is None:
print(f"Error reading {mask_name} or {raw_name}. Skipping...")
continue

# Create an empty weight map with the same dimensions as the raw image
weight_map = np.zeros_like(raw)
print(raw.shape)
print(mask.shape)

# Set pixels where mask == 38 to the color [0, 0, 255] (blue)
weight_map[mask == 38] = [0, 0, 100]
# Create an empty weight map with the same dimensions as the raw image
weight_map = np.zeros_like(raw)

# Overlay the weight map onto the original image
overlay = cv2.addWeighted(raw, 1, weight_map, 0.5, 0)

# Save the intermediate results
# cv2.imwrite(os.path.join(output_dir, f"{raw_name}_mask.jpg"), mask)
# # cv2.imwrite(os.path.join(output_dir, f"{raw_name}_binary_mask.jpg"), binary_mask)
# cv2.imwrite(os.path.join(output_dir, f"{raw_name}_weight_map.jpg"), weight_map)
# Set pixels where mask == 38 to the color [0, 0, 255] (blue)
weight_map[mask == 38] = [0, 0, 100]

# Overlay the weight map onto the original image
overlay = cv2.addWeighted(raw, 1, weight_map, 0.5, 0)

# Save the intermediate results
# cv2.imwrite(os.path.join(output_dir, f"{raw_name}_mask.jpg"), mask)
# # cv2.imwrite(os.path.join(output_dir, f"{raw_name}_binary_mask.jpg"), binary_mask)
# cv2.imwrite(os.path.join(output_dir, f"{raw_name}_weight_map.jpg"), weight_map)

# Save the overlaid image to the output directory
output_name = os.path.join(output_dir, raw_name)
cv2.imwrite(output_name, overlay)
print(f"Saved {output_name}")

print("Processing complete!")
# Save the overlaid image to the output directory
output_name = os.path.join(output_dir, raw_name)
cv2.imwrite(output_name, overlay)
print(f"Saved {output_name}")

print(f"Processed {mask_name}")

# Process stainmasks
print("Processing stainmasks...")
process_masks(stain_mask_dir, stain_output_dir)

# Process filteredStainMasks
print("Processing filteredStainMasks...")
process_masks(filtered_mask_dir, filtered_output_dir)

print("All processing complete!")
5 changes: 1 addition & 4 deletions stainsegmentation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,10 +32,7 @@
os.makedirs(mask_directory)

# Get the model path
# model_path = Path("/home/roboticslab/Developer/pytorch_concrete_flaws_segmentation") / model
# print(f"MODEL_PATH={model_path}")

model_path = "/image2overlays/best_model_stain.pth"
model_path = config["StainSegmentation"]["model"]
print(f"MODEL_PATH={model_path}")

# Define the directory where inference.py is located
Expand Down

0 comments on commit 9197073

Please sign in to comment.