from stereoimage_generation import create_stereoimages
from tqdm import tqdm
import argparse
from pathlib import Path
from PIL import Image
import numpy as np
import os


def invert_depth_map(depth_map):
    inverted_depth_map = depth_map * -1
    return inverted_depth_map


def get_matching_depth_map(image_file, depth_map_files):
    base_image_name = Path(image_file).stem
    matching_depth_maps = [f for f in depth_map_files if base_image_name in str(f)]
    if matching_depth_maps:
        return matching_depth_maps[0]
    else:
        print(f"Warning: No matching depth map found for {base_image_name}.")
        return None


def create_stereo_images(input_image, depth_map, output_path, divergence, separation, save_right_only):
    try:
        inverted_depth_map = invert_depth_map(depth_map)
        stereo_images = create_stereoimages(input_image, inverted_depth_map, divergence=divergence,
                                            separation=separation, modes=['left-right'],
                                            fill_technique='polylines_sharp')

        if save_right_only:
            width, height = stereo_images[0].size
            right_half = stereo_images[0].crop((width // 2, 0, width, height))
            output_image = right_half
        else:
            output_image = stereo_images[0]

        output_path_jpg = output_path.with_suffix('.jpg')
        output_image.save(output_path_jpg, format='JPEG', quality=100)

    except AssertionError as e:
        print(f"Ошибка при обработке файла {output_path}: {e}")


def process_images(images_dir, depth_maps_dir, output_dir, divergence, separation, save_right_only, delete_depth_maps):
    image_files = sorted(Path(images_dir).glob("*.jpg"))
    depth_map_files = sorted(Path(depth_maps_dir).glob("*.png"))

    output_dir_path = Path(output_dir)
    output_dir_path.mkdir(parents=True, exist_ok=True)

    for image_file in tqdm(image_files, desc="Обработка изображений"):
        matching_depth_map = get_matching_depth_map(image_file, depth_map_files)

        if matching_depth_map is None:
            continue

        try:
            input_image = Image.open(image_file)
            depth_map = np.array(Image.open(matching_depth_map))
        except (IOError, ValueError) as e:
            print(f"Ошибка при открытии файла {image_file} или соответствующей карты глубины: {e}")
            continue

        output_path = output_dir_path / image_file.name

        if not output_path.with_suffix('.jpg').exists():
            create_stereo_images(input_image, depth_map, output_path, divergence, separation, save_right_only)

            # Delete the depth map if the flag is set
            if delete_depth_maps:
                try:
                    os.remove(matching_depth_map)
                except Exception as e:
                    print(f"Не удалось удалить карту глубины {matching_depth_map}: {e}")


def main():
    parser = argparse.ArgumentParser(description="Создание стереоизображений из изображений и карт глубины.")
    parser.add_argument("--images_dir", required=True, help="Путь к директории с оригинальными изображениями")
    parser.add_argument("--depth_maps_dir", required=True, help="Путь к директории с картами глубины")
    parser.add_argument("--output_dir", required=True, help="Выходной путь для сохранения стереоизображений")
    parser.add_argument("--divergence", type=float, default=2.5,
                        help="Степень расхождения между левым и правым изображениями (по умолчанию: 2.5)")
    parser.add_argument("--separation", type=int, default=0,
                        help="Расстояние между левым и правым изображениями (по умолчанию: 0)")
    parser.add_argument("--save_right_only", action='store_true', help="Сохранять только правую половину изображения")
    parser.add_argument("--delete_depth_maps", action='store_true',
                        help="Удалять обработанные карты глубины после создания стереоизображений")

    args = parser.parse_args()

    process_images(args.images_dir, args.depth_maps_dir, args.output_dir, args.divergence, args.separation,
                   args.save_right_only, args.delete_depth_maps)


if __name__ == "__main__":
    main()


Мне надо изменить флаг save_right_only. Нужно его сделать так, что бы при его передаче сохранять отдельно левое и правое изображение, с префиксами _L и _R