data_filter.py

import os 
import torch
import numpy as np
import cv2
data_path="data/task_0003/train"
cams=['038522063145','037522062165']
all_demos = sorted(os.listdir(data_path))
all_demos = [demo for demo in all_demos if "task" in demo]

dry_run = False   # Change to False if you want to actually delete

for demo in all_demos:
    print("Processing: ",demo)
    first_camera_color_dir = os.path.join(data_path,demo, f"cam_{cams[0]}", 'color')
    print(first_camera_color_dir)
    # input()
    first_camera_depth_dir = os.path.join(data_path,demo, f"cam_{cams[0]}", 'depth')
    first_file_names = [f for f in os.listdir(first_camera_color_dir) 
                        if os.path.isfile(os.path.join(first_camera_color_dir, f))
                        and f.endswith(".png") ]
    first_timestamps = [int(name.replace('.png', '')) for name in first_file_names]
    ##################################################################################################
    tcp_dir=os.path.join(data_path,demo,"lowdim")
    tcp=np.load(f"{tcp_dir}/tcp.npz",allow_pickle=False)
    # Filter the timestamps, and get the new first_timestamps, use os.remove for removing the filtered images in the first camera dir, both color and depth
    timestamp=np.array(sorted(tcp.files,key=int))
    # clear_tcp = np.stack([tcp[k] for k in timestamp], axis=0)  # (T, 7)
    
    first_idx=0
    for ts in timestamp:
        if not ((tcp[ts][3:])**2).sum()==0:
            first_idx=ts
            break
    print("first valid timestamp:",first_idx)
    timestamps=np.array([t for t in timestamp if t >= first_idx])
    tcp = np.stack([tcp[k] for k in timestamps], axis=0)
    filtered_timestamp=[timestamps[0]]
    
    for i in range(1,len(timestamps)):
        trans1,trans2=tcp[i-1,:3],tcp[i,:3]
        rot1,rot2=tcp[i-1,3:],tcp[i,3:]
        trans=np.linalg.norm(trans2-trans1)
        dot = np.dot(rot1, rot2)
        dot = np.clip(abs(dot), -1.0, 1.0) 
        rot = 2 * np.arccos(dot)
        moved=(trans>1e-4) or (rot>0)
        if moved:
            print("trans:",trans)
            print(rot)
            filtered_timestamp.append(timestamps[i])
        else:
            print("unmoved timestamp:",timestamps[i])
            print("trans:",trans)
            print("rot:",rot)
        
    first_timestamps=[int(ft) for ft in filtered_timestamp]
    # Read the first frame to get size
    first_frame = cv2.imread(os.path.join(data_path, demo, f"cam_{cams[0]}", 'color', str(first_timestamps[0]))+".png")
    height, width, _ = first_frame.shape

    # Define video writer
    out_path = "output_video.mp4"
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # codec
    fps = 30  # frames per second
    out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
    i=0
    print(len(first_timestamps))
    for ts in first_timestamps:
        i=i+1
        image_path=os.path.join(data_path, demo, f"cam_{cams[0]}", 'color', str(ts))+".png"
        print(i,image_path)
        frame = cv2.imread(image_path)
        if frame is None:
            print(f"Warning: could not read {image_path}")
            continue
        out.write(frame)
    import pdb;pdb.set_trace()
    out.release()
    print(f"Video saved to {out_path}")
    
    first_to_delete=[f for f in first_file_names if int(f.replace('.png','')) not in first_timestamps]
    print(len(first_to_delete))
    for f in first_to_delete:
        file_path = os.path.join(first_camera_color_dir, f)
        os.remove(file_path)
        file_path = os.path.join(first_camera_depth_dir, f)
        os.remove(file_path)
    # import pdb;pdb.set_trace()
    ##################################################################################################
    second_camera_color_dir = os.path.join(data_path, demo, f"cam_{cams[1]}", 'color')
    second_camera_depth_dir = os.path.join(data_path, demo, f"cam_{cams[1]}", 'depth')
    second_file_names = [f for f in os.listdir(second_camera_color_dir) 
                         if os.path.isfile(os.path.join(second_camera_color_dir, f))]
    second_timestamps = [int(name.replace('.png', '')) for name in second_file_names]

    # Build pairs
    pairs = {}
    for ts in first_timestamps:
        closest = min(second_timestamps, key=lambda x: abs(x - ts))
        pairs[str(ts)] = closest
    torch.save(pairs,f"{first_camera_color_dir}/pairs.pth")
    # Camera2 frames that are paired
    matched_second = set(pairs.values())

    # Camera2 frames that should be deleted
    to_delete = [f for f in second_file_names if int(f.replace('.png','')) not in matched_second]

    # print(f"Demo: {demo}")
    if dry_run:
        print("These files would be deleted from camera2:")
        print(to_delete)
    else:
        # print("Deleting files from camera2:")
        for f in to_delete:
            file_path = os.path.join(second_camera_color_dir, f)
            os.remove(file_path)
            file_path = os.path.join(second_camera_depth_dir, f)
            os.remove(file_path)
            # print(f"Deleted {file_path}")