Added a setup.py file for a local raft package

demo.py

@@ -1,6 +1,3 @@
-import sys
-sys.path.append('core')
-
 import argparse
 import os
 import cv2
@@ -10,9 +7,8 @@
 from PIL import Image
 
 from raft import RAFT
-from utils import flow_viz
-from utils.utils import InputPadder
-
+from raft.utils import flow_viz
+from raft.utils.utils import InputPadder
 
 
 DEVICE = 'cuda'
@@ -41,7 +37,7 @@ def viz(img, flo):
 
 def demo(args):
     model = torch.nn.DataParallel(RAFT(args))
-    model.load_state_dict(torch.load(args.model))
+    model.load_state_dict(torch.load(args.model, map_location=DEVICE))
 
     model = model.module
     model.to(DEVICE)

evaluate.py

@@ -1,6 +1,3 @@
-import sys
-sys.path.append('core')
-
 from PIL import Image
 import argparse
 import os
@@ -10,12 +7,12 @@
 import torch.nn.functional as F
 import matplotlib.pyplot as plt
 
-import datasets
-from utils import flow_viz
-from utils import frame_utils
+import raft.datasets as datasets
+from raft.utils import flow_viz
+from raft.utils import frame_utils
 
 from raft import RAFT
-from utils.utils import InputPadder, forward_interpolate
+from raft.utils.utils import InputPadder, forward_interpolate
 
 
 @torch.no_grad()

raft/__init__.py

@@ -0,0 +1 @@
+from raft.raft import RAFT

core/corr.py → raft/corr.py

@@ -1,6 +1,6 @@
 import torch
 import torch.nn.functional as F
-from utils.utils import bilinear_sampler, coords_grid
+from raft.utils.utils import bilinear_sampler, coords_grid
 
 try:
     import alt_cuda_corr

core/datasets.py → raft/datasets.py

@@ -11,8 +11,8 @@
 from glob import glob
 import os.path as osp
 
-from utils import frame_utils
-from utils.augmentor import FlowAugmentor, SparseFlowAugmentor
+from raft.utils import frame_utils
+from raft.utils.augmentor import FlowAugmentor, SparseFlowAugmentor
 
 
 class FlowDataset(data.Dataset):

core/raft.py → raft/raft.py

@@ -3,10 +3,10 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from update import BasicUpdateBlock, SmallUpdateBlock
-from extractor import BasicEncoder, SmallEncoder
-from corr import CorrBlock, AlternateCorrBlock
-from utils.utils import bilinear_sampler, coords_grid, upflow8
+from raft.update import BasicUpdateBlock, SmallUpdateBlock
+from raft.extractor import BasicEncoder, SmallEncoder
+from raft.corr import CorrBlock, AlternateCorrBlock
+from raft.utils.utils import bilinear_sampler, coords_grid, upflow8
 
 try:
     autocast = torch.cuda.amp.autocast

setup.py

@@ -0,0 +1,16 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='raft',
+    version='1.0.0',
+    packages=find_packages(),
+    install_requires=[
+        'torch',
+        'torchvision',
+        'opencv-python',
+        'matplotlib',
+        'tensorboard',
+        'scipy'
+    ]
+)
+

train.py

@@ -1,6 +1,4 @@
 from __future__ import print_function, division
-import sys
-sys.path.append('core')
 
 import argparse
 import os
@@ -16,8 +14,8 @@
 
 from torch.utils.data import DataLoader
 from raft import RAFT
+import raft.datasets as datasets
 import evaluate
-import datasets
 
 from torch.utils.tensorboard import SummaryWriter
 
@@ -244,4 +242,4 @@ def train(args):
     if not os.path.isdir('checkpoints'):
         os.mkdir('checkpoints')
 
-    train(args)
+    train(args)