File size: 8,645 Bytes
5b3b0f4 |
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 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 |
import cv2
import numpy as np
import os
from multiprocessing import Pool
def pad_to_match_height(image, target_height, color=(255, 255, 255)):
h, w, c = image.shape
if h >= target_height:
return image # No padding needed
top_pad = 0
bottom_pad = target_height - h - top_pad
return cv2.copyMakeBorder(image, top_pad, bottom_pad, 0, 0, cv2.BORDER_CONSTANT, value=color)
def compress(img, resize_factor, path):
if image_path.find("kitti")!=-1:
return img
if image_path.find("eth3d")!=-1:
return img
if image_path.find("middlebury")!=-1:
return np.hstack([img[:, 100:650, :],
img[:, 850:1400, :],
img[:, 1600:2150, :]])
elif image_path.find("booster")!=-1:
return np.hstack([img[:, int(590 * resize_factor):int(2550 * resize_factor), :],
img[:, int(3500 * resize_factor):int(5550 * resize_factor), :],
img[:, int(6450 * resize_factor):int(8450 * resize_factor), :]])
else:
raise Exception("Not supported: ", path)
def split_image_with_custom_grouping(image_path, output_dir, num_rows, resize_factor=None, compression_params=None, gap=10, destination_path=None):
# Read the image
img = cv2.imread(image_path)
if img is None:
raise ValueError(f"Cannot read the image from {image_path}")
if resize_factor is not None:
img = cv2.resize(img, (0, 0), fx=resize_factor, fy=resize_factor, interpolation=cv2.INTER_AREA)
# cv2.imwrite(os.path.join(output_dir, "resized_image.png"), img)
# Get the height and width of the image
height, width, _ = img.shape
# Determine the height of each row
row_height = height // num_rows
# Create the output directory
os.makedirs(output_dir, exist_ok=True)
# Combine the first five rows
first_five_height = 5 * row_height
first_five_rows = img[:first_five_height, :, :]
if image_path.find("kitti")!=-1:
group = []
for i in range(0, 5):
upper = i * row_height + 50
lower = (i + 1) * row_height - 40
group.append(img[upper:lower, :, :])
first_five_rows = np.vstack(group)
compressed_first_five_rows = compress(first_five_rows, resize_factor, image_path)
# cv2.imwrite(os.path.join(output_dir, "first_five_rows.jpg"), compressed_first_five_rows, compression_params)
# Initialize groups
group_0 = [] # Rows divisible by 3
group_1 = [] # Rows with remainder 1 when divided by 3
group_2 = [] # Rows with remainder 2 when divided by 3
group_3 = [] # Additional rows with remainder 2
upper_shift, lower_shit = 0, 0
if image_path.find("kitti")!=-1:
upper_shift = 50
lower_shit = -40
# Process the remaining rows
for i in range(5, num_rows):
if i < 29:
upper = i * row_height + upper_shift
lower = (i + 1) * row_height + lower_shit if i < num_rows - 1 else height
else:
upper = 29 * row_height + (i - 29) * (row_height - 1) + upper_shift
lower = 29 * row_height + (i - 28) * (row_height - 1) + lower_shit if i < num_rows - 1 else height
row_img = img[upper:lower, :, :]
compressed_row_img = compress(row_img, resize_factor, image_path)
if i < num_rows - 2:
# Group rows based on modulo 3
if (i - 5) % 3 == 0:
group_0.append(compressed_row_img)
elif (i - 5) % 3 == 1:
group_1.append(compressed_row_img)
else:
group_2.append(compressed_row_img)
else:
group_3.append(compressed_row_img)
# Concatenate and save each group
if group_0:
group_0_img = np.vstack(group_0) # Concatenate vertically
# cv2.imwrite(os.path.join(output_dir, "group_0.jpg"), group_0_img, compression_params)
if group_1:
group_1_img = np.vstack(group_1)
# cv2.imwrite(os.path.join(output_dir, "group_1.jpg"), group_1_img, compression_params)
if group_2:
group_2_img = np.vstack(group_2)
# cv2.imwrite(os.path.join(output_dir, "group_2.jpg"), group_2_img, compression_params)
if group_3:
group_3_img = np.vstack(group_3)
# cv2.imwrite(os.path.join(output_dir, "group_3.jpg"), group_3_img, compression_params)
if destination_path is not None:
group_images = []
group_images.append(np.vstack(group_1))
group_images.append(np.vstack(group_0))
group_images.append(np.vstack(group_2))
# Create the gap as a black (zeros) image
gap_img = np.ones((gap, group_3_img.shape[1], 3), dtype=np.uint8) * 255
left_img = np.vstack([compressed_first_five_rows, gap_img, group_3_img])
# Find the maximum height among all images
max_height = max(left_img.shape[0], *(img.shape[0] for img in group_images))
# Pad the first image to match the maximum height
padded_image1 = pad_to_match_height(left_img, max_height)
padded_images = [pad_to_match_height(img, max_height) for img in group_images]
gap_img = np.ones((max_height, gap, 3), dtype=np.uint8) * 255
concatenated_img = padded_image1
for img in padded_images:
concatenated_img = np.hstack([concatenated_img, gap_img, img])
# Save the result
cv2.imwrite(destination_path.replace("png","jpg"), concatenated_img, compression_params)
def process_file(args):
source_path, output_dir, num_rows, resize_factor, compression_params, gap, destination_path = args
split_image_with_custom_grouping(source_path, output_dir, num_rows, resize_factor, compression_params, gap, destination_path)
print(f"Processed {source_path} to {destination_path}")
def process(input_dir, output_dir, num_rows, resize_factor=None, compression_params=None, gap = 10):
# Ensure the output directory exists
os.makedirs(output_dir, exist_ok=True)
tasks = []
# Walk through the directory tree
for root, dirs, files in os.walk(input_dir):
for file in files:
# Check if the file is an image based on extension
if file.lower().endswith(('.png',)):
# Determine the relative path of the current directory
relative_path = os.path.relpath(root, input_dir)
# Create the corresponding output subdirectory
output_subdir = os.path.join(output_dir, relative_path)
os.makedirs(output_subdir, exist_ok=True)
# Copy the file to the output directory
source_path = os.path.join(root, file)
destination_path = os.path.join(output_subdir, file.replace(".png", ".jpg"))
tasks.append((source_path, output_dir, num_rows, resize_factor, compression_params, gap, destination_path))
# split_image_with_custom_grouping(source_path, output_dir, num_rows, resize_factor, compression_params, gap, destination_path)
# print(f"process {source_path} to {destination_path}")
with Pool(processes=(os.cpu_count())//2) as pool: # Use as many processes as available CPU cores
pool.map(process_file, tasks)
image_path = "/data5/yao/runs/vis/inter_results/RaftStereoDepthBetaK53DispRefineSigmoidPreMonoBatch48ConfDim_20241102_014050/SUPP/analysis/booster/train-balanced-Case-disp_00.png"
# image_path = "/data5/yao/runs/vis/inter_results/RaftStereoDepthBetaK53DispRefineSigmoidPreMonoBatch48ConfDim_20241102_014050/SUPP/analysis/eth3d/two_view_training-delivery_area_1l-disp0GT.png"
# image_path = "/data5/yao/runs/vis/inter_results/RaftStereoDepthBetaK53DispRefineSigmoidPreMonoBatch48ConfDim_20241102_014050/SUPP/analysis/kitti/training-disp_occ_0-000000_10.png"
# image_path = "/data5/yao/runs/vis/inter_results/RaftStereoDepthBetaK53DispRefineSigmoidPreMonoBatch48ConfDim_20241102_014050/SUPP/analysis/middlebury_h/MiddEval3-trainingH-Adirondack-disp0GT.png"
output_dir = "./output"
num_rows = 43
resize_factor = 0.25
compression_params = [cv2.IMWRITE_JPEG_QUALITY, 75]
destination_path = "./output/res.png"
# split_image_with_custom_grouping(image_path, output_dir, num_rows, resize_factor, compression_params, 10, destination_path)
input_dir = "/data5/yao/runs/vis/inter_results/RaftStereoDepthBetaK53DispRefineSigmoidPreMonoBatch48ConfDim_20241102_014050/SUPP/analysis/"
output_dir = "./output"
process(input_dir, output_dir, num_rows, resize_factor, compression_params, gap=10)
|