Removing flow quantization from optical flow scripts. Fixing README.md.

Author: ponimatkin

README.md

@@ -38,7 +38,6 @@ conda env create --name ssl-vos python=3.8 pip
 conda install pytorch torchvision torchaudio cudatoolkit=11.6 -c pytorch -c conda-forge
 pip install scikit-learn opencv-python fast_slic imageio matplotlib scikit-image easydict path.py
 
-git clone https://github.com/ylabbe/bullet3.git && cd bullet3 
 python setup.py build
 python setup.py install
 ```
@@ -55,7 +54,7 @@ in order to prepare rest of your environment.
 Our approach requires three steps to run produce segmentations. At first, extract DINO features and optical flows by via
 ```
 python extract_dino_features.py --dataset davis
-python generate_flow_arflow.py --dataset davis --step 1 --raw
+python generate_flow_arflow.py --dataset davis --step 1
 ```
 
 The second step consists of extracting the initial eigenvectors, which can be done by

generate_flow_arflow.py

@@ -10,9 +10,9 @@
 
 from ARFlow.utils.torch_utils import restore_model
 from ARFlow.models.pwclite import PWCLite
-from config import DATA_PATH, MAX_FLOW_VAL
-from utils.flow_utils import quantize_flow
 from ARFlow.utils.warp_utils import flow_warp
+from config import DATA_PATH
+
 
 
 def calculate_binary_flow_weights(org_img, warped_img, thr_value, thr_type):
@@ -87,12 +87,7 @@ def generate_arflow_flow(args):
 
                 binary_weights = calculate_binary_flow_weights(image1, re_image1, args.thr_value, args.thr_type)
                 flow = flow * binary_weights[:,:,None]
-                if args.raw:
-                    np.save(flow_folder / imfile1.with_suffix('.npy').name, flow)
-                else:
-                    dx, dy = quantize_flow(flow, max_val=MAX_FLOW_VAL, norm=False)
-                    flow = np.stack([dx, dy, np.zeros(dx.shape)], axis=-1)
-                    cv2.imwrite((flow_folder / imfile1.name).as_posix(), flow[:, :, [2, 1, 0]])
+                np.save(flow_folder / imfile1.with_suffix('.npy').name, flow)
 
             images = list(reversed(images))
             print(f'Working on folder: {folder.name} in backward direction')
@@ -111,20 +106,14 @@ def generate_arflow_flow(args):
 
                 binary_weights = calculate_binary_flow_weights(image1, re_image1, args.thr_value, args.thr_type)
                 flow = flow * binary_weights[:,:,None]
-                if args.raw:
-                    np.save(flow_folder_reverse / imfile1.with_suffix('.npy').name, flow)
-                else:
-                    dx, dy = quantize_flow(flow, max_val=MAX_FLOW_VAL, norm=False)
-                    flow = np.stack([dx, dy, np.zeros(dx.shape)], axis=-1)
-                    cv2.imwrite((flow_folder_reverse / imfile1.name).as_posix(), flow[:, :, [2, 1, 0]])
+                np.save(flow_folder_reverse / imfile1.with_suffix('.npy').name, flow)
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--model', help="restore checkpoint", default='ARFlow/checkpoints/CityScapes/pwclite_ar.tar')
     parser.add_argument('--dataset', help="dataset for flow estimation")
     parser.add_argument('--step', type=int, default=1, help="flow step size")
     parser.add_argument('--test-shape', default=[384, 640], type=int, nargs=2)
-    parser.add_argument('--raw', action='store_true', help='generate raw optical flow')
     parser.add_argument('--thr_type', type=str, default='percentile')
     parser.add_argument('--thr_value', type=int, default=90)
     args = parser.parse_args()

generate_flow_raft.py

@@ -8,8 +8,7 @@
 from RAFT.core.raft import RAFT
 from RAFT.core.utils.utils import InputPadder
 
-from config import DATA_PATH, MAX_FLOW_VAL
-from utils.flow_utils import quantize_flow
+from config import DATA_PATH
 
 
 def load_image(imfile):
@@ -59,12 +58,7 @@ def generate_raft_flow(args):
 
                 _, flow = model(image1, image2, iters=20, test_mode=True)
                 flow = flow[0].permute(1, 2, 0).cpu().numpy()
-                if args.raw:
-                    np.save(flow_folder / imfile1.with_suffix('.npy').name, flow)
-                else:
-                    dx, dy = quantize_flow(flow, max_val=MAX_FLOW_VAL, norm=False)
-                    flow = np.stack([dx, dy, np.zeros(dx.shape)], axis=-1)
-                    cv2.imwrite((flow_folder / imfile1.name).as_posix(), flow[:, :, [2, 1, 0]])
+                np.save(flow_folder / imfile1.with_suffix('.npy').name, flow)
 
             images = list(reversed(images))
             print(f'Working on folder: {folder.name} in backward direction')
@@ -77,20 +71,14 @@ def generate_raft_flow(args):
 
                 _, flow = model(image1, image2, iters=20, test_mode=True)
                 flow = flow[0].permute(1, 2, 0).cpu().numpy()
-                if args.raw:
-                    np.save(flow_folder_reverse / imfile1.with_suffix('.npy').name, flow)
-                else:
-                    dx, dy = quantize_flow(flow, max_val=MAX_FLOW_VAL, norm=False)
-                    flow = np.stack([dx, dy, np.zeros(dx.shape)], axis=-1)
-                    cv2.imwrite((flow_folder_reverse / imfile1.name).as_posix(), flow[:, :, [2, 1, 0]])
+                np.save(flow_folder_reverse / imfile1.with_suffix('.npy').name, flow)
 
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--model', help="restore checkpoint")
     parser.add_argument('--dataset', help="dataset for flow estimation")
     parser.add_argument('--step', type=int, default=1, help="flow step size")
-    parser.add_argument('--raw', action='store_true', help='generate raw optical flow')
     parser.add_argument('--small', action='store_true', help='use small model')
     parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision')
     parser.add_argument('--alternate_corr', action='store_true', help='use efficent correlation implementation')