tools/split_vis_image.py

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)