Locally Installable Library for RAFT-Stereo

demo.py

@@ -1,14 +1,11 @@
-import sys
-sys.path.append('core')
-
 import argparse
 import glob
 import numpy as np
 import torch
 from tqdm import tqdm
 from pathlib import Path
-from raft_stereo import RAFTStereo
-from utils.utils import InputPadder
+from raft_stereo.raft_stereo import RAFTStereo
+from raft_stereo.utils.utils import InputPadder
 from PIL import Image
 from matplotlib import pyplot as plt
 

evaluate_stereo.py

@@ -1,16 +1,14 @@
 from __future__ import print_function, division
-import sys
-sys.path.append('core')
 
 import argparse
 import time
 import logging
 import numpy as np
 import torch
 from tqdm import tqdm
-from raft_stereo import RAFTStereo, autocast
-import stereo_datasets as datasets
-from utils.utils import InputPadder
+from raft_stereo.raft_stereo import RAFTStereo, autocast
+import raft_stereo.stereo_datasets as datasets
+from raft_stereo.utils.utils import InputPadder
 
 def count_parameters(model):
     return sum(p.numel() for p in model.parameters() if p.requires_grad)

raft_stereo/__init__.py

@@ -0,0 +1 @@
+from raft_stereo import RAFTStereo

core/corr.py → raft_stereo/corr.py

@@ -1,6 +1,6 @@
 import torch
 import torch.nn.functional as F
-from core.utils.utils import bilinear_sampler
+from raft_stereo.utils.utils import bilinear_sampler
 
 try:
     import corr_sampler

core/raft_stereo.py → raft_stereo/raft_stereo.py

@@ -1,10 +1,10 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from core.update import BasicMultiUpdateBlock
-from core.extractor import BasicEncoder, MultiBasicEncoder, ResidualBlock
-from core.corr import CorrBlock1D, PytorchAlternateCorrBlock1D, CorrBlockFast1D, AlternateCorrBlock
-from core.utils.utils import coords_grid, upflow8
+from raft_stereo.update import BasicMultiUpdateBlock
+from raft_stereo.extractor import BasicEncoder, MultiBasicEncoder, ResidualBlock
+from raft_stereo.corr import CorrBlock1D, PytorchAlternateCorrBlock1D, CorrBlockFast1D, AlternateCorrBlock
+from raft_stereo.utils.utils import coords_grid, upflow8
 
 
 try:

core/stereo_datasets.py → raft_stereo/stereo_datasets.py

@@ -14,8 +14,8 @@
 from glob import glob
 import os.path as osp
 
-from core.utils import frame_utils
-from core.utils.augmentor import FlowAugmentor, SparseFlowAugmentor
+from raft_stereo.utils import frame_utils
+from raft_stereo.utils.augmentor import FlowAugmentor, SparseFlowAugmentor
 
 
 class StereoDataset(data.Dataset):

setup.py

@@ -0,0 +1,22 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='raft_stereo',
+    version='1.0.0',
+    packages=find_packages(),
+    install_requires=[
+        'torch',
+        'torchvision',
+        'opencv-python',
+        'matplotlib',
+        'tensorboard',
+        'scipy',
+        'opt_einsum',
+        'glob',
+        'skimage',
+        'Pillow',
+        'imageio',
+        'numpy'
+    ]
+)
+

train_stereo.py

@@ -10,10 +10,10 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from core.raft_stereo import RAFTStereo
+from raft_stereo.raft_stereo import RAFTStereo
 
 from evaluate_stereo import *
-import core.stereo_datasets as datasets
+import raft_stereo.stereo_datasets as datasets
 
 try:
     from torch.cuda.amp import GradScaler