curation_pipeline/select_frame_with_this_that.py

'''
    This repository is used to prepare Bridge dataset with this that conditioning
'''
import os, sys, shutil
import pickle
from ultralytics import YOLO
from PIL import Image, ImageDraw
import numpy as np
import cv2
import math
import collections
from segment_anything import SamAutomaticMaskGenerator, SamPredictor, sam_model_registry


def show_mask(mask, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)

    return mask_image * 255


def read_center_point(model, img_path, do_visualization, store_path):

    action_img = Image.open(img_path)
    prediction = model.predict(source=action_img, save=False)[0]  # Only 1 frame
    
    if not hasattr(prediction, "boxes"):
        print("Detection Fail: We cannot have boxes attribute")
        return None, None   # -1 means NAN and pass this case
    
    # save at the temp_places for visualizaiton
    if do_visualization:
        prediction.save(filename=store_path)


    bounding_boxes = prediction.boxes.xywh
    num, dim = bounding_boxes.shape
    assert(dim == 4)
    
    # Catch up all center point of all bounding boxes
    edge_point_cord = []
    center_points = []
    for idx in range(num):
        x, y, w, h = bounding_boxes[idx].detach().cpu().numpy()
        center_point = [x, y]   # TODO: y+(h/4) 根据经验，往下飘逸25%的高度，一般来说比较有帮助

        edge_point_cord.extend([ (x+w//2, y+h//2), (x-w//2, y+h//2), (x-w//2, y-h//2), (x+w//2, y-h//2) ])


        if w <= 15 or h <= 15:    # If a bounding box is too small, we will disregard this case
            return None, None

        # Calculate the distance between current one and previous points for sanity check
        for point in center_points: # Check all previous points
            give_up_threshold = 90
            if center_point[0] - point[0] >= give_up_threshold:  
                print("Two points are too far away and neglect the case")
                return None, None
            if center_point[1] - point[1] >= give_up_threshold:
                print("Two points are too far away and neglect the case")
                return None, None
        
        # Append to the list
        center_points.append(center_point)

    
    if len(center_points) == 0 or len(center_points) > 2:
        print("Detection Fail: We cannot detect bounding boxes")
        return None, None
    
    # Calculating the average distance among center_points
    if len(center_points) == 2:
        first_box, second_box = center_points
        
        center_x = (first_box[0] + second_box[0]) / 2
        center_y = (first_box[1] + second_box[1]) / 2

        distance = math.sqrt(abs(first_box[0] - second_box[0])**2 + abs(first_box[1] - second_box[1])**2)
        
        return [center_x, center_y, distance], edge_point_cord
        
    return [*center_points[0], 100], edge_point_cord # if len(center_points) == 1, distance is 0; however, to avoid 2-1-2 box detection in sequential, we set it as a higher value



def detect_gripper(gripper_detection_model, input_dir, action_start, action_end, do_visualization, store_dir, sample_failure_collect_folder=None):

    # 先处理第一个point的（这个比较重要，所以要重复3次）；然后再快速处理最后一个point

    # Process the first action frame by iterating next three frames and choose the closest one
    first_center_points = []
    edge_point_cords = []
    for idx in range(3):    # Repeat 3 times
        action_start_path = os.path.join(input_dir, "im_"+str(action_start + idx)+".jpg")
        first_center_point, edge_point_cord = read_center_point(gripper_detection_model, action_start_path, do_visualization, os.path.join(store_dir, "contact_first"+str(idx)+".jpg"))    # The first frame
        
        if idx == 0 and first_center_point is None:
            message = "Cannot find the first contact point!"

            print("The contact point we cannot detect is at ", action_start_path)
            if sample_failure_collect_folder != "":
                shutil.copyfile(action_start_path, os.path.join(sample_failure_collect_folder, str(len(os.listdir(sample_failure_collect_folder)))+".jpg") )

            return (None, None, message)

        if first_center_point is not None:
            first_center_points.append([action_start + idx, first_center_point])

            # Add edge points
            print(edge_point_cord)
            edge_point_cords.extend(edge_point_cord)    # 我有点担心所有point就这么extend会对一些的edge case不是那么robust


    # Select the closest point between two
    first_center_points.sort(key=lambda x: x[1][2])
    first_center_point = first_center_points[0][1][:2]
    start_idx = first_center_points[0][0]
    print("first_center_point is " + str(first_center_point)  + " with idx " + str(start_idx))
    order_idx = [start_idx, action_end]
    

    # Find the xmin, ymin, xmax, ymax for based all three points as the bounding box for the SAM
    edge_point_cords.sort(key=lambda x: x[0])
    xmin = int(edge_point_cords[0][0])
    xmax = int(edge_point_cords[-1][0])

    edge_point_cords.sort(key=lambda x: x[1])
    ymin = int(edge_point_cords[0][1])
    ymax = int(edge_point_cords[-1][1])

    bbox_info = (xmin, xmax, ymin, ymax)


    # Process the last action frame
    action_end_path = os.path.join(input_dir, "im_"+str(action_end)+".jpg")
    last_center_point, edge_point_cord = read_center_point(gripper_detection_model, action_end_path, do_visualization, os.path.join(store_dir, "contact_last.jpg"))  # The last frame
    if last_center_point is None:
        message = "Cannot find the last contact point!"

        print("The contact point we cannot detect is at ", action_start_path)
        if sample_failure_collect_folder != "":
            store_name = str(len(os.listdir(sample_failure_collect_folder))) + ".jpg"
            shutil.copyfile(action_start_path, os.path.join(sample_failure_collect_folder, store_name) )

        return (None, bbox_info, message)
    last_center_point = last_center_point[:2]
    

    # Check if two center points is too close, if they are too close, we will merge to one point
    merge_threshold = 30
    if math.sqrt((first_center_point[0] - last_center_point[0])**2 + (first_center_point[1] - last_center_point[1])**2) <= merge_threshold:
        print("Merge two points to one!")
        message = "Success!"
        return ([[first_center_point], order_idx], bbox_info, message)


    # Return needed information
    message = "Success!"
    return ([[first_center_point, last_center_point], order_idx], bbox_info, message)




def visualize_this_that(base_img, bbox_info, this_that_points):

    # Draw a green dot only for the start point
    for point in this_that_points:
        print("point is ", point)
        target_horizontal, target_vertical = point
        target_horizontal, target_vertical = int(target_horizontal), int(target_vertical)
        
        dot_range = 3 
        for i in range(-1*dot_range, dot_range+1):
            for j in range(-1*dot_range, dot_range+1):
                dil_vertical, dil_horizontal = target_vertical + i, target_horizontal + j
                if (0 <= dil_vertical and dil_vertical < base_img.shape[0]) and (0 <= dil_horizontal and dil_horizontal < base_img.shape[1]):
                    base_img[dil_vertical, dil_horizontal, :] = [0, 128, 0]
                # else:
                #     # print("The traj is out of boundary!!!!!!!!!!!!!!!!!!!!! and we won't consider it")      # 现在
                #     return (False, base_img)
            
    # Draw the bounding box
    xmin, xmax, ymin, ymax = bbox_info
    base_img = cv2.rectangle(base_img, (xmin, ymin), (xmax, ymax), color=(0,0,255), thickness=2)

    return (True, base_img)



def manage_seq_range(input_dir, store_dir, sample_failure_collect_folder, total_frames_needed, 
                        max_original_input_tolerate, gripper_detection_model, sam_predictor, do_visualization):

    # Find valid image lists
    num_frames_input = 0
    for file_name in os.listdir(input_dir):
        if file_name.startswith("im_"):
            num_frames_input += 1
    for idx in range(num_frames_input):
        target_path = os.path.join(input_dir, "im_"+str(idx)+".jpg")
        if not os.path.exists(target_path):
            print("We don't have ", target_path)
            message = "Invalid error"   # Make sure that every file in this order is existed, this is quite important
            return (False, message)
    

    if num_frames_input > max_original_input_tolerate:
        message = "The number of frames is too long for constructing the sequence length needed"
        return (False, message)
    
    if num_frames_input < total_frames_needed:
        message = "The number of frames is too short for constructing the sequence length needed"
        return (False, message)
    


    # Prepare this and that based on policy_out.pkl
    policy_out_file_path = os.path.join(input_dir, "policy_out.pkl")
    with open(policy_out_file_path, "rb") as f:
        policy = pickle.load(f)

    actions_codes = []
    action_start, action_end = None, None
    for idx, item in enumerate(policy):
        action_value = item["actions"][-1]
        if action_start is None and action_value == 0.0:
            action_start = idx

        if (action_start is not None) and (action_end is None) and (action_value == 1.0): 
            action_end = idx    # Until record the first 1.0 exists after the first 0.0 appears
        actions_codes.append(action_value)
        
    if action_start is None or action_end is None:  
        message = "We cannot read an action_start or action_end code!"
        return (False, message)    # Requires to have both start and end actions (Usually, they are a pair)
    
    print("actions_codes is ", actions_codes)
    print("the start end idx we read is ", action_start, action_end)


    # Detect the gripper (should return a list with exactly two x,y coordinate points)
    detection_retrun_info, bbox_info, detect_message = detect_gripper(
                                                                        gripper_detection_model, 
                                                                        input_dir,
                                                                        action_start,
                                                                        action_end,
                                                                        do_visualization = do_visualization,
                                                                        store_dir = store_dir,
                                                                        sample_failure_collect_folder = sample_failure_collect_folder, 
                                                                    )
    if detection_retrun_info is None:
        return (False, detect_message)
    
    detected_point, old_seq_idx = detection_retrun_info
    print("detected_point is ", detected_point)


    # Visualize if needed    
    base_img = cv2.imread(os.path.join(input_dir, "im_0.jpg"))
    if do_visualization:
        status, visual_img = visualize_this_that(base_img, bbox_info, detected_point)
        if status:
            cv2.imwrite(os.path.join(store_dir, "visualization.png"), visual_img)



    # SAM process based on bbox_info
    xmin, xmax, ymin, ymax = bbox_info
    sam_predictor.set_image(np.uint8(base_img))
    positive_point_cords = np.array([[ int(detected_point[0][0]), int(detected_point[0][1]) ]])
    positive_point_cords = np.array(positive_point_cords)
    positive_point_labels = np.ones(len(positive_point_cords))

    # Predict the mask based on the point and bounding box designed
    masks, scores, logits = sam_predictor.predict(  
                                                    point_coords = positive_point_cords,
                                                    point_labels = positive_point_labels,
                                                    box = np.array([xmin, ymin, xmax, ymax])[None, :],
                                                    multimask_output = False,
                                                )
    print(scores)
    for mask_idx, mask in enumerate(masks):
        mask_img = show_mask(mask)
        cv2.imwrite(os.path.join(store_dir, "mask_" + str(mask_idx) + ".png"), mask_img)



    ################################ Move the img ######################################
    # Calculate needed parameters
    division_factor = num_frames_input // total_frames_needed
    remain_frames = (num_frames_input % total_frames_needed) - 1    # -1 for adaptation

    # Define the gap
    gaps = [division_factor for _ in range(total_frames_needed-1)]
    for idx in range(remain_frames):
        if idx % 2 == 0:
            gaps[idx//2] += 1      # Start to end order
        else:
            gaps[-1*(1+(idx//2))] += 1   # End to start order

    # Map the gap to the specific orders
    idx_orders = [1]    # 从1还是shift一下问题应该不大
    for global_idx, gap in enumerate(gaps):
        idx_orders.append(idx_orders[-1] + gap)
    if idx_orders[-1] >= num_frames_input:
        message = "Invalid error"
        return (False, message)
    # assert(idx_orders[-1] < num_frames_input)
    assert(len(idx_orders) == total_frames_needed)
    

    # Copy the essential files first
    for global_idx, cur_idx in enumerate(idx_orders):
        source_path = os.path.join(input_dir, "im_"+str(cur_idx)+".jpg")
        destination_path = os.path.join(store_dir, "im_"+str(global_idx)+".jpg")

        if not os.path.exists(source_path):     # Theoretically, source_path must exists
            message = "We couldn't find the source path. Theoretically, source_path must exists!"  # 有一种可能就是我们丢失了一些地方，在cp或者本来就没有，记得统计数量
            return (False, message)

        shutil.copyfile(source_path, destination_path)

    # Map order_idx to the cropped version
    mapped_seq_idx = []
    for old_idx in old_seq_idx:
        tmp = []
        for tmp_idx, new_idx in enumerate(range(len(idx_orders))):
            tmp.append((tmp_idx, abs(old_idx - idx_orders[new_idx])))
        # Sort the smallest fistance
        tmp.sort(key=lambda x: x[1])
        mapped_seq_idx.append(tmp[0][0])

    print("Before the idx is ", old_seq_idx)
    print("mapped idx is ", mapped_seq_idx)


    # Write the information to new destination
    f = open(os.path.join(store_dir, "data.txt"), "a")
    f.write(str(mapped_seq_idx[0]) + " " + str(detected_point[0][0]) + " " + str(detected_point[0][1]) + "\n")
    if len(detected_point) == 2:   # Two points excluding the last idx 
        f.write(str(mapped_seq_idx[1]) + " " + str(detected_point[1][0]) + " " + str(detected_point[1][1]) + "\n")
    f.close()


    # Move lang.txt file
    shutil.copyfile(os.path.join(input_dir, 'lang.txt'), os.path.join(store_dir, 'lang.txt'))


    message = "Success!"
    return (True, message)




if __name__ == "__main__":

    # General storage setting
    dataset_path = "../datasets_rob/Bridge_v2_raw"
    destination_path = "../sanity_check/bridge_v2_TT14_longer_tolerance"
    sample_failure_collect_folder = ""      # This is to collect cases that fail for active learning

    total_frames_needed = 14
    max_original_input_tolerate = 56        # 40 for 14 fps; 60 for 25fps; 
    do_visualization = True


    # YOLO model init
    yolo_pretarined_path = "pretrained/yolov8n_best.pt"
    gripper_detection_model = YOLO("yolov8n.yaml")  # build a new model from scratch
    gripper_detection_model = YOLO(yolo_pretarined_path)  # load a pretrained model (recommended for training)

    # SAM model init
    model_type = "vit_h"
    sam_pretrained_path = "pretrained/sam_vit_h_4b8939.pth"
    sam = sam_model_registry[model_type](checkpoint=sam_pretrained_path).to(device="cuda")
    sam_predictor = SamPredictor(sam)     # There is a lot of setting here


    # Make dir if needed
    if os.path.exists(destination_path):
        shutil.rmtree(destination_path)    
    os.makedirs(destination_path)

    # Prepare the folder to collect failure cases
    if sample_failure_collect_folder != "":
        if os.path.exists(sample_failure_collect_folder):
            shutil.rmtree(sample_failure_collect_folder) 
        os.makedirs(sample_failure_collect_folder)



    # Collect the message
    message_dict = collections.defaultdict(int)


    store_idx = 0
    for folder_name in sorted(os.listdir(dataset_path)):
        input_folder_path = os.path.join(dataset_path, folder_name)
        store_folder_path = os.path.join(destination_path, "0"*(6-len(str(store_idx)))+str(store_idx))
        print("We are processing ", input_folder_path)

        # Prepare store_folder_path folder
        os.makedirs(store_folder_path)
        
        status, message = manage_seq_range(input_folder_path, store_folder_path, sample_failure_collect_folder, total_frames_needed, max_original_input_tolerate, gripper_detection_model, sam_predictor, do_visualization)
        if status:      # We will only update the store_idx only when this file is successfully written
            store_idx += 1
        else:
            print("This status failed! Message: " + message)
            shutil.rmtree(store_folder_path)
        # break # For debug
            
        # Collect the infor to dict
        message_dict[message] += 1
    
    print("We have " + str(store_idx) + " valid dataset")
    print("message_dict info is ", message_dict)