Untitled

import os
import numpy as np
import json
from PIL import Image
# import wandb
from glob import glob
from moviepy.editor import ImageSequenceClip, CompositeVideoClip, TextClip, ImageClip, concatenate_videoclips, clips_array
from moviepy.config import get_setting
from moviepy.config import change_settings

print("xxxx", get_setting("IMAGEMAGICK_BINARY"))
IMAGEMAGICK_BINARY = os.getenv('IMAGEMAGICK_BINARY', 'auto-detect')
print("xxxx", get_setting("IMAGEMAGICK_BINARY"))

# os.system("convert example_images/monkey.jpeg output.jpg")

# stroke_color = "green"
# stroke_width = 3

position = ('left', 'bottom')  # Bottom-left corner
font_size = 48
text = "Text"
color = 'white'
stroke_color = 'black'
stroke_width = 2

# texts = ["Hello", "World", "MoviePy"]
# sizes = [24, 36, 48]  # Different font sizes for each line
# colors = ["red", "green", "blue"]  # Different colors for each line
# positions = [("center", "top"), ("center", "center"), ("center", "bottom")]  # Different positions

def find_step(number):
    number += 1
    if 1 <= number <= 7:
        return f"{1}, {number}"
    elif 8 <= number <= 17:
        return f"{2}, {number - 7}"
    elif 18 <= number <= 27:
        return f"{3}, {number - 17}"
    elif 28 <= number <= 37:
        return f"{4}, {number - 27}"
    else:
        return ""

def create_complex_layout_textclip(texts, styles, duration):
    clips = []
    for text, style in zip(texts, styles):
        # Create each TextClip with its specific style
        txt_clip = TextClip(text, fontsize=style['fontsize'], color=style['color'], font=style['font'], size=(640, 480))
        txt_clip = txt_clip.set_position(style['position']).set_duration(duration)
        clips.append(txt_clip)
    
    return clips
    
def add_text_to_image(image_path, additional_image_path, meta_data, text, position, font_size, color):
    # print(meta_data = {meta_data["original_prompt"], 
    #                    meta_data["editing_prompt"]})

    h, w = 256, 256
    black_image_array = ImageClip(np.zeros((h, w, 3), dtype=np.uint8),  duration=1)
    black_image_array_bt = ImageClip(np.zeros((int(h/4), w, 3), dtype=np.uint8),  duration=1)
    img_clip_input = ImageClip("../../data/PIE-Bench_v1/annotation_images/" + meta_data["image_path"], duration=1).resize((h, w))  # duration in seconds per frame

    # Load the image as a clip
    img_clip = ImageClip(image_path, duration=1).resize((h, w))  # duration in seconds per frame
    basename = os.path.basename(image_path)
    dirname = os.path.dirname(image_path)
    edt_path = glob(os.path.join(dirname, '_'.join(basename.replace("Rec", "Edit").split('_')[:-1]) + '**.png'))
    # print(edt_path)
    # import pdb; pdb.set_trace()
    assert len(edt_path) == 1
    img_clip_edt = ImageClip(edt_path[0], duration=1).resize((h, w))  # duration in seconds per frame

    # Load the additional image
    additional_img_clip = ImageClip(additional_image_path).resize((h, w)).set_duration(1)

    add_basename = os.path.basename(additional_image_path)
    add_dirname = os.path.dirname(additional_image_path)
    add_edt_path = glob(os.path.join(add_dirname, '_'.join(add_basename.replace("Rec", "Edit").split('_')[:-1]) + '**.png'))
    assert len(add_edt_path) == 1
    additional_img_clip_edt = ImageClip(add_edt_path[0]).resize((h, w)).set_duration(1)

    # # Resize additional image to match main image height if necessary
    # additional_img_clip = additional_img_clip

    # Concatenate the two images horizontally
    # print(img_clip_input.shape, img_clip.shape, black_image_array.shape, img_clip_edt.shape)
    concatenated_clip = clips_array([[img_clip_input, img_clip, additional_img_clip], 
                                     [black_image_array, img_clip_edt, additional_img_clip_edt],
                                     [black_image_array_bt,black_image_array_bt,black_image_array_bt]])



    y_pos = 30
    # Create a text clip
    txt_clip1 = TextClip('[Edit:] '+meta_data["editing_prompt"], fontsize=font_size/2, color=color, font='Arial')
    txt_clip1 = txt_clip1.set_position(("left", 2*h + h/4 - y_pos)).set_duration(1)

    y_pos += 24
    txt_clip2 = TextClip('[Rec:] '+meta_data["original_prompt"], fontsize=font_size/2, color=color, font='Arial')
    txt_clip2 = txt_clip2.set_position(("left",  2*h + h/4 - y_pos)).set_duration(1)

    y_pos += 50
    txt_clip = TextClip(text, fontsize=font_size, color=color, font='Arial')
    txt_clip = txt_clip.set_position(("left", 2*h + h/4 - y_pos)).set_duration(1)



    # Composite the text clip over the image clip
    composite_clip = CompositeVideoClip([concatenated_clip, txt_clip1, txt_clip2, txt_clip])
    return composite_clip

# def add_text_to_frame(frame_array, text, position, font_size, color, duration=1 / 30):
#     # Convert the frame array to a MoviePy ImageClip
#     frame_clip = ImageClip(frame_array).set_duration(duration)

#     # Create a TextClipq
#     text_clip = TextClip(text, fontsize=font_size, color=color)
#     text_clip = text_clip.set_position(position).set_duration(duration)

#     # Create a CompositeVideoClip combining the frame and text
#     composite_clip = CompositeVideoClip([frame_clip, text_clip])

#     return composite_clip

run_id_reno1 = "0excf4a5"
run_id_reno2 = "cir0jluw"
run_id_rnri = "ffwbgbjd"

all_runs = glob('wandb/**')
run_name_dict = {os.path.basename(fx).split('-')[-1]: fx for fx in all_runs if not os.path.isfile(fx)}

imgage_id = '000000000003'
tag = 'Rec'
all_image_items = set(sorted([os.path.basename(xt).split('_')[0].split(' ')[0] for xt in glob(os.path.join(run_name_dict[run_id_rnri], f'files/media/images/**{tag}*.png'))]))
# print()
# import pdb; pdb.set_trace()

annotation_file_name = os.path.join("../../data/PIE-Bench_v1/mapping_file.json")
with open (annotation_file_name) as f:
    annotation_file = json.load(f)


for imgage_id in all_image_items:
    
    rnri_files = glob(os.path.join(run_name_dict[run_id_rnri], f'files/media/images/{imgage_id} {tag}*.png'))
    reno_files =  glob(os.path.join(run_name_dict[run_id_reno1], f'files/media/images/{imgage_id} {tag}*.png'))
    reno_files.extend(glob(os.path.join(run_name_dict[run_id_reno2], f'files/media/images/{imgage_id} {tag}*.png')))
    

    rnri_files = sorted(rnri_files)
    reno_files = sorted(reno_files)

    if len(rnri_files) != len(reno_files):
        print("skip, as:", len(rnri_files), "!==", len(reno_files))
        continue

    base_rnri = int(os.path.basename(rnri_files[0]).split('_')[1]) 
    base_reno = int(os.path.basename(reno_files[0]).split('_')[1])


    # Example usage:
    os.makedirs('output_video', exist_ok=True)
    output_video_file = f'output_video/{imgage_id}.mp4'
    frame_rate = 30  # frames per second



    # # Create the video clip
    # clip = ImageSequenceClip(rnri_files, fps=frame_rate)

    # Create a list to hold frames with text
    frames_with_text = []

    # for frame in rnri_files:
    # Load the frame
    # Add text to the frame
    # idx = int(os.path.basename(frame).split('_')[1]) - base_rnri
    # Create a list of clips with text added
    clips = [add_text_to_image(img, img2, annotation_file[imgage_id], find_step(int(os.path.basename(img).split('_')[1]) - base_reno), 
                            position, font_size, color) for img, img2 in zip(reno_files, rnri_files)]

    # Concatenate all the clips together
    final_clip = concatenate_videoclips(clips, method="compose")

    # Write the result to a file
    final_clip.write_videofile(output_video_file, fps=frame_rate, codec='libx264')  # fps can be adjusted


    for img_name_rnri, img_name_reno in zip(rnri_files, reno_files):
        # print(os.path.basename(img_name1), os.path.basename(img_name2))
        Image.open(img_name_rnri).convert("RGB")
        Image.open(img_name_reno).convert("RGB")