Function single_process

preprocess/track_FrameINO.py:405–741

(csv_folder_path, store_folder_path, GPU_offset, debug)

@torch.no_grad()
def single_process(csv_folder_path, store_folder_path, GPU_offset, debug):
    
    # Setting
    store_freq = 10
    device = 'cuda'


    # Read the csv file
    csv_idx = GPU_offset
    csv_file_path = os.path.join(csv_folder_path, "sub" + str(csv_idx) + ".csv")
    print("CSV file we read is ", csv_file_path)


    # Prepare the store file path
    store_file_path = os.path.join(store_folder_path, "sub" + str(csv_idx) + ".csv")
    if not resume and os.path.exists(store_file_path):
        # Remove existing csv
        os.remove(store_file_path)


    # Resume the store csv  
    find_resume = True
    if resume:      # Read the last store row
        find_resume = False
        with open(store_file_path, 'r') as file:
            reader = csv.reader(file)
            store_rows = list(reader)
            last_store_row = store_rows[-1]
            print("The number of rows we have processed in the store csv is ", len(store_rows))
            


    # Init the SAM2 model
    sam2_predictor = SAM2VideoPredictor.from_pretrained("facebook/sam2.1-hiera-large")     # Large model has 224.4M Param at 39.5 FPS



    # Read all row in the csv file
    start_time = time.time()
    info_lists = []       # The order will be follow automatically
    with open(csv_file_path) as file_obj:

        reader_obj = csv.reader(file_obj) 
        
        # Iterate over each row in the csv  
        for row_idx, row in enumerate(reader_obj): 

            # For the first row case (With all title content)
            if row_idx == 0:    # The first line is the title of content
                elements = dict()
                for element_idx, key in enumerate(row):
                    elements[key] = element_idx

                print("The first row is ", row + ["ID_info"])       # TODO: we need ID_path (Especially 1st frame) + Region BBox position (All corresponds to Motion Traj)

                # Store the csv
                if not resume:
                    with open(store_file_path, 'a', newline='') as csvfile: