import sys import argparse import copy import os, shutil import imageio import cv2 from PIL import Image, ImageDraw import os.path as osp import random import numpy as np import torch.multiprocessing as mp from multiprocessing import set_start_method import math, time, gc import torch import torch.nn.functional as F import matplotlib.pyplot as plt from segment_anything import SamAutomaticMaskGenerator, SamPredictor, sam_model_registry # Import files from the local path root_path = os.path.abspath('.') sys.path.append(root_path) from config.flowformer_config import get_cfg from flowformer_code.utils import flow_viz, frame_utils from flowformer_code.utils.utils import InputPadder from flowformer_code.FlowFormer import build_flowformer TRAIN_SIZE = [432, 960] def show_anns(anns): if len(anns) == 0: return sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True) ax = plt.gca() ax.set_autoscale_on(True) img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4)) img[:,:,3] = 0 for ann in sorted_anns: m = ann['segmentation'] color_mask = np.concatenate([np.random.random(3), [0.35]]) img[m] = color_mask return img*255 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 compute_grid_indices(image_shape, patch_size=TRAIN_SIZE, min_overlap=20): if min_overlap >= TRAIN_SIZE[0] or min_overlap >= TRAIN_SIZE[1]: raise ValueError( f"Overlap should be less than size of patch (got {min_overlap}" f"for patch size {patch_size}).") if image_shape[0] == TRAIN_SIZE[0]: hs = list(range(0, image_shape[0], TRAIN_SIZE[0])) else: hs = list(range(0, image_shape[0], TRAIN_SIZE[0] - min_overlap)) if image_shape[1] == TRAIN_SIZE[1]: ws = list(range(0, image_shape[1], TRAIN_SIZE[1])) else: ws = list(range(0, image_shape[1], TRAIN_SIZE[1] - min_overlap)) # Make sure the final patch is flush with the image boundary hs[-1] = image_shape[0] - patch_size[0] ws[-1] = image_shape[1] - patch_size[1] return [(h, w) for h in hs for w in ws] def compute_flow(model, image1, image2, weights=None): print(f"computing flow...") image_size = image1.shape[1:] image1, image2 = image1[None].cuda(), image2[None].cuda() hws = compute_grid_indices(image_size) if weights is None: # no tile padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_pre, _ = model(image1, image2) flow_pre = padder.unpad(flow_pre) flow = flow_pre[0].permute(1, 2, 0).cpu().numpy() else: # tile flows = 0 flow_count = 0 for idx, (h, w) in enumerate(hws): image1_tile = image1[:, :, h:h+TRAIN_SIZE[0], w:w+TRAIN_SIZE[1]] image2_tile = image2[:, :, h:h+TRAIN_SIZE[0], w:w+TRAIN_SIZE[1]] flow_pre, _ = model(image1_tile, image2_tile) padding = (w, image_size[1]-w-TRAIN_SIZE[1], h, image_size[0]-h-TRAIN_SIZE[0], 0, 0) flows += F.pad(flow_pre * weights[idx], padding) flow_count += F.pad(weights[idx], padding) flow_pre = flows / flow_count flow = flow_pre[0].permute(1, 2, 0).cpu().numpy() return flow def compute_adaptive_image_size(image_size): target_size = TRAIN_SIZE scale0 = target_size[0] / image_size[0] scale1 = target_size[1] / image_size[1] if scale0 > scale1: scale = scale0 else: scale = scale1 image_size = (int(image_size[1] * scale), int(image_size[0] * scale)) return image_size def prepare_image(viz_root_dir, fn1, fn2, keep_size): print(f"preparing image...") image1 = frame_utils.read_gen(fn1) image2 = frame_utils.read_gen(fn2) image1 = np.array(image1).astype(np.uint8)[..., :3] image2 = np.array(image2).astype(np.uint8)[..., :3] if not keep_size: dsize = compute_adaptive_image_size(image1.shape[0:2]) image1 = cv2.resize(image1, dsize=dsize, interpolation=cv2.INTER_CUBIC) image2 = cv2.resize(image2, dsize=dsize, interpolation=cv2.INTER_CUBIC) image1 = torch.from_numpy(image1).permute(2, 0, 1).float() image2 = torch.from_numpy(image2).permute(2, 0, 1).float() dirname = osp.dirname(fn1) filename = osp.splitext(osp.basename(fn1))[0] viz_dir = osp.join(viz_root_dir, dirname) # if not osp.exists(viz_dir): # os.makedirs(viz_dir) viz_fn = osp.join(viz_dir, filename + '.png') return image1, image2, viz_fn def build_model(): print(f"building model...") cfg = get_cfg() model = torch.nn.DataParallel(build_flowformer(cfg)) model.load_state_dict(torch.load(cfg.model)) model.cuda() model.eval() return model def filter_uv(flow, threshold_factor = 0.2): u = flow[:,:,0] v = flow[:,:,1] rad = np.sqrt(np.square(u) + np.square(v)) rad_max = np.max(rad) threshold = threshold_factor * rad_max flow[:,:,0][rad < threshold] = 0 flow[:,:,1][rad < threshold] = 0 return flow def visualize_traj(base_img, traj_path, connect_points = True): target_vertical, target_horizontal = traj_path[-1] if connect_points and len(traj_path) > 1: # Draw a line to connect two point to show motion direction start_coordinate = (traj_path[-2][1], traj_path[-2][0]) end_coordinate = (traj_path[-1][1], traj_path[-1][0]) pil_img = Image.fromarray(base_img) # Draw the line color = 'red' draw = ImageDraw.Draw(pil_img) draw.line([start_coordinate, end_coordinate], fill = color, width = 3) base_img = np.array(pil_img) # Draw a green dot only for the start point if len(traj_path) == 1: 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) return (True, base_img) def calculate_flow(viz_root_dir, store_dir, img_pairs, optical_flow_model, sam_predictor, SAM_positive_sample_num, SAM_negative_sample_num, mask_generator, traj_visualization, keep_size, verbose=False): # Trajectory prepare traj_path = [] # It collects all points traversed in a temporal order is_hard_to_track = False # If this is True, it means that, we have a time in tracking hard to find dx and dy movement. Under this circumstance, we are not very recommended to use it hard_track_idxs = set() traj_image_lists = [] # Iterate all image pairs for idx, img_pair in enumerate(img_pairs): fn1, fn2 = img_pair print(f"processing {fn1}, {fn2}...") image1, image2, viz_fn = prepare_image(viz_root_dir, fn1, fn2, keep_size) # Be very careful, image1 and image2 may be different resolution shape if keep_size is False # Generate the optical flow and filter those that is small motion flow_uv = filter_uv(compute_flow(optical_flow_model, image1, image2, None)) # if verbose: # Store the visualization of flow_uv # flow_img = flow_viz.flow_to_image(flow_uv) # cv2.imwrite("optical_flow_" + str(idx+1) + ".png", flow_img[:, :, [2,1,0]]) if idx == 0: # We will store the first image to memory for further visualization purpose # Base img # base_img = np.uint8(np.transpose(image1.numpy(), (1,2,0))) # SAM figure # sam_all = mask_generator.generate(image1) # base_img = show_anns(sam_all) # base_img = np.transpose(base_img, (1,2,0)) # Plain white image base_img = np.zeros(np.transpose(image1.numpy(), (1,2,0)).shape, dtype=np.uint8) base_img.fill(255) # Extract moving points (positive point) positive_point_cords = [] nonzeros = np.nonzero(flow_uv) # [(vertical), (horizontal)] if len(nonzeros[0]) < SAM_positive_sample_num: # We require the number of points to be more than SAM_positive_sample_num return False positive_orders = np.random.choice(len(nonzeros[0]), SAM_positive_sample_num, replace=False) # we have randomly select instead of use all in the sam_predictor prediction for i in range(len(nonzeros[0])): if i in positive_orders: positive_point_cords.append([nonzeros[1][i], nonzeros[0][i]]) # 根据document来看，这个就应该是先horizontal再vertical，也就是这个顺序 positive_point_cords = np.array(positive_point_cords) positive_point_labels = np.ones(len(positive_point_cords)) # Define negative sample (outside the optical flow choice) if SAM_negative_sample_num != 0: skip_prob = 2 * SAM_negative_sample_num / (flow_uv.shape[0]*flow_uv.shape[1] - len(nonzeros[0])) negative_point_cords = [] for i in range(flow_uv.shape[0]): for j in range(flow_uv.shape[1]): if flow_uv[i][j][0] == 0 and flow_uv[i][j][1] == 0: # 0 means the no motion zone and we have already filter low motion as zero before if random.random() < skip_prob: negative_point_cords.append([j, i]) # 根据document来看，这个就应该是先horizontal再vertical，也就是这个顺序 negative_point_cords = np.array(negative_point_cords) # [:SAM_negative_sample_num] negative_point_labels = np.zeros(len(negative_point_cords)) # Make sure that it is less than / equals to SAM_negative_sample_num quantity ################## Use SAM to filter out what we need (& use negative points) ################## if idx == 0: # Only consider the first frame now. # With sample coordinate sam_predictor.set_image(np.uint8(np.transpose(image1.numpy(), (1,2,0)))) if SAM_negative_sample_num != 0 and len(negative_point_cords) != 0: all_point_cords = np.concatenate((positive_point_cords, negative_point_cords), axis=0) all_point_labels = np.concatenate((positive_point_labels, negative_point_labels), axis=0) else: all_point_cords = positive_point_cords all_point_labels = positive_point_labels masks, scores, logits = sam_predictor.predict( point_coords=all_point_cords, point_labels=all_point_labels, multimask_output=False, ) mask = masks[0] # TODO: 一定要确定我们这里选择了最大的mask，而没有考虑的第二大和其他的, 这里可能有bug，我们默认了第一个就是最大的mask # if verbose: # cv2.imwrite("mask_"+str(idx+1)+".png", (np.uint8(mask)*255)) # annotated_img = show_mask(mask) # cv2.imwrite("annotated.png", annotated_img) ################## Choose the one we need as the reference for the future tracking ################## # Choose a random point in the mask target_zone = np.nonzero(mask) # [(vertical), (horizontal)] target_zone = [(target_zone[0][i], target_zone[1][i]) for i in range(len(target_zone[0]))] # Now, the sturcture is [(vertical, horizontal), ...] repeat_time = 0 loop2find = True while loop2find: loop2find = False start_point = target_zone[np.random.choice(len(target_zone), 1, replace=False)[0]] start_vertical, start_horizontal = start_point repeat_time += 1 if repeat_time == 100: # In some minor case, it may have infinite loop, so we need to manually break if it is looping print("We are still hard to find a optimal first point, but we cannot let it loop") break # Try to choose a start_point that is more centralized (Not close to the border) fast_break = False for i in range(-15, 15): for j in range(-15, 15): dil_vertical, dil_horizontal = start_vertical + i, start_horizontal + j if (0 <= dil_vertical and dil_vertical < mask.shape[0]) and (0 <= dil_horizontal and dil_horizontal < mask.shape[1]): if mask[dil_vertical][dil_horizontal] == 0: print("We need to change to a new position for the start p Since this one is close to the border of the object...........") loop2find = True fast_break = True break else: # We won't want to consider those that is close to the boundary print("We need to change to a new position Since this one is close to the border of the image...........") loop2find = True fast_break = True break if fast_break: break traj_path.append(start_point) status, base_img = visualize_traj(base_img, traj_path) if status == False: # If the traj is False, we won't consider it anymore. file = open("log.txt", "a") file.write("Invalid start point\n") return False # Read from the last one in traj ref_vertical, ref_horizontal = traj_path[-1][0], traj_path[-1][1] # Get the average motion vector for point surrounding (8+1 directions) the ref_point; This is because this is the most accurate statistics horizon_lists, vertical_lists = [], [] start_range, end_range = -5, 5 # Calculate the average motion based on surrounding motion search_times = 0 while len(horizon_lists) == 0: # If we cannot find a direction, we use average value inside this mask, but we will flag it. search_times += 1 if search_times > 1: print("This is hard to track!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! and we have tracked " + str(search_times) + " times") # TODO: 如果out of boundary那种，search times到了8-10次的就砍掉那后面frame吧，这种非常inaccurate了， 你也可以retrack一个新的点，但是没有什么意义，看整体数量来定吧 is_hard_to_track = True hard_track_idxs.add(idx) if abs(start_range) >= flow_uv.shape[0]//2: file = open("log.txt", "a") file.write("This folder has search all space but didn't find any place to track optical flow\n") return False # If we have already search for the whole graph but didn't find anything to track, we discard this sample # Search for a larger space which is nearby 我觉得扩大搜索范围应该是最稳定的选择吧 for i in range(start_range, end_range): for j in range(start_range, end_range): target_vertical, target_horizontal = ref_vertical + i, ref_horizontal + j if 0 <= target_vertical and target_vertical < flow_uv.shape[0] and 0 <= target_horizontal and target_horizontal < flow_uv.shape[1]: if flow_uv[target_vertical, target_horizontal, 0] == 0 or flow_uv[target_vertical, target_horizontal, 1] == 0: continue # Ignore zero vector to ensure only calculate moving position horizon_lists.append(flow_uv[target_vertical, target_horizontal, 0]) # Horizontal motion strength vertical_lists.append(flow_uv[target_vertical, target_horizontal, 1]) # Vertical motion strength # If there isn't any to search, we kepp on a larger space start_range -= 10 end_range += 10 average_dx = sum(horizon_lists)/len(horizon_lists) average_dy = sum(vertical_lists)/len(vertical_lists) print("average movement is ", (average_dx, average_dy)) traj_path.append(( int(traj_path[-1][0] + average_dy), int(traj_path[-1][1] + average_dx))) # Append the motion in independent order print(traj_path) ##################### Visualize the trajectory path (Debug Purpose) ##################### status, base_img = visualize_traj(base_img, traj_path) if status == False: # If the traj is False, we won't consider it anymore. return False cv2.imwrite(os.path.join(store_dir, "traj_path.png"), cv2.cvtColor(base_img, cv2.COLOR_BGR2RGB)) if traj_visualization: status, single_traj_img = visualize_traj(np.uint8(np.transpose(image1.numpy(), (1,2,0))), traj_path[:-1], connect_points=False) if status == False: # If the traj is False, we won't consider it anymore. return False traj_write_path = os.path.join(store_dir, "traj_"+str(idx)+".png") # cv2.imwrite(traj_write_path, cv2.cvtColor(single_traj_img, cv2.COLOR_BGR2RGB)) traj_image_lists.append(traj_write_path) # if traj_visualization: # images = [] # for filename in traj_image_lists: # images.append(imageio.imread(filename)) # # os.remove(filename) # Remove when used # imageio.mimsave(os.path.join(store_dir, 'traj_motion.gif'), images, duration=0.05) # TODO: 可以如果hard to track，就aggressivly多试即便，我们根据这个hard_track_idxs的长度来粗略判断哪个最好，三次里面选最好的 if is_hard_to_track: if len(hard_track_idxs) >= len(img_pairs)//3: # If more than half of the traj is hard to track, we need to consider discard this one file = open("log.txt", "a") file.write("we have a lot of times hard to find dx and dy movement. Under this circumstance, we are not very recommended to use the track\n") return False # Write a file store all position for further utilization txt_path = os.path.join(store_dir, "traj_data.txt") if os.path.exists(txt_path): os.remove(txt_path) file = open(txt_path, "a") for traj in traj_path: file.write(str(traj[0]) + " " + str(traj[1]) + "\n") # Save in numpy information # with open(os.path.join(store_dir, 'traj_data.npy'), 'wb') as f: # np.save(f, flow_uv) print("We write ", traj_path) return True def manage_seq_range(input_dir, store_dir, total_frame_needed): lists = os.listdir(input_dir) lists = lists[2:-2] num_frames_input = len(lists) if num_frames_input < total_frame_needed: print("The number of frames is too short for constructing the sequnece length needed") return False division_factor = num_frames_input // total_frame_needed remain_frame = num_frames_input % total_frame_needed gaps = [division_factor for _ in range(total_frame_needed)] for idx in range(remain_frame): gaps[idx] += 1 cur_idx = 2 for global_idx, gap in enumerate(gaps): source_path = os.path.join(input_dir, "im_"+str(cur_idx)+".jpg") destination_path = os.path.join(store_dir, "im_"+str(global_idx)+".jpg") shutil.copyfile(source_path, destination_path) cur_idx += gap return True def generate_pairs(dirname, start_idx, end_idx): img_pairs = [] for idx in range(start_idx, end_idx): img1 = osp.join(dirname, f'im_{idx}.jpg') img2 = osp.join(dirname, f'im_{idx+1}.jpg') # img1 = f'{idx:06}.png' # img2 = f'{idx+1:06}.png' img_pairs.append((img1, img2)) return img_pairs def process_partial_request(request_list, num_frames, traj_visualization, viz_root_dir): # Init the optical flow model optical_flow_model = build_model() # Init SAM for segmentation task model_type = "vit_h" weight_path = "pretrained/sam_vit_h_4b8939.pth" SAM_positive_sample_num = 20 # How many points we use for the positive sample num () SAM_negative_sample_num = 0 # How many points we use for the negative sample num print("In multi processing, we will build an instance of mask_generator independently") sam = sam_model_registry[model_type](checkpoint=weight_path).to(device="cuda") mask_generator = SamAutomaticMaskGenerator(sam) print("In multi processing, we will build an instance of sam_predictor independently") sam_predictor = SamPredictor(sam) counter = 0 while True: counter += 1 if counter == 10: counter = 0 gc.collect() print("We will sleep here to clear memory") time.sleep(5) info = request_list[0] request_list = request_list[1:] if info == None: print("This queue ends") break # Process each sub_input_dir and store the information there sub_input_dir = info img_pairs = generate_pairs(sub_input_dir, 0, num_frames-1) print(img_pairs) with torch.no_grad(): # Calculate the optical flow and return a status to say whther this generated flow is usable status = calculate_flow(viz_root_dir, sub_input_dir, img_pairs, optical_flow_model, sam_predictor, SAM_positive_sample_num, SAM_negative_sample_num, mask_generator, traj_visualization, keep_size = True) # file = open("log.txt", "a") print("The status for folder " + sub_input_dir + " is " + str(status) + "\n") if status == False: # If the status is failed, we will remove it afterwords print("The status is Failed, so we won't store this one as one promising data") else: print("We have successfully process one!") if __name__ == '__main__': # Manage the paramter parser = argparse.ArgumentParser() parser.add_argument('--input_dir', default = '../validation_flow14/') parser.add_argument('--num_workers', type = int, default = 1) # starting index of the image sequence parser.add_argument('--viz_root_dir', default = 'viz_results') parser.add_argument('--traj_visualization', default = True) # If this is True, # list_start = 0 # list_end = 25000 num_frames = 14 args = parser.parse_args() input_dir = args.input_dir num_workers = args.num_workers viz_root_dir = args.viz_root_dir traj_visualization = args.traj_visualization store_idx = 0 dir_list = [] for sub_input_name in sorted(os.listdir(input_dir)): sub_input_dir = os.path.join(input_dir, sub_input_name) # sub_store_dir = os.path.join(store_dir, "0"*(7-len(str(store_idx)))+str(store_idx)) store_idx += 1 dir_list.append(sub_input_dir) # Truncate the list to the target # dir_list = dir_list[list_start:] # Use multiprocessing to handle to speed up num = math.ceil(len(dir_list) / num_workers) for idx in range(num_workers): # set_start_method('spawn', force=True) request_list = dir_list[:num] request_list.append(None) dir_list = dir_list[num:] process_partial_request(request_list, num_frames, traj_visualization, viz_root_dir) # This is for debug purpose # p = mp.Process(target=process_partial_request, args=(request_list, num_frames, traj_visualization, viz_root_dir, )) # p.start() print("Submitted all jobs!") # p.join() # 好像不加这个multiprocess就莫名自己结束了 print("All task finished!")