import time
from tensorboardX import SummaryWriter
from datasets import __datasets__
from models import __models__
from utils import *
from torch.utils.data import DataLoader
import gc
import skimage
cudnn.benchmark = True
parser = argparse.ArgumentParser(description='Cascade Stereo Network (CasStereoNet)')
parser.add_argument('--model', default='gwcnet-c', help='select a model structure', choices=__models__.keys())
parser.add_argument('--maxdisp', type=int, default=192, help='maximum disparity')
parser.add_argument('--test_dataset', required=True, help='dataset name', choices=__datasets__.keys())
parser.add_argument('--test_datapath', required=True, help='data path')
parser.add_argument('--testlist', required=True, help='testing list')
parser.add_argument('--test_batch_size', type=int, default=1, help='testing batch size')
parser.add_argument('--logdir', required=True, help='the directory to save logs and checkpoints')
parser.add_argument('--loadckpt', help='load the weights from a specific checkpoint')
parser.add_argument("--local_rank", type=int, default=0)
parser.add_argument('--ndisps', type=str, default="48,24", help='ndisps')
parser.add_argument('--disp_inter_r', type=str, default="4,1", help='disp_intervals_ratio')
parser.add_argument('--dlossw', type=str, default="0.5,2.0", help='depth loss weight for different stage')
parser.add_argument('--cr_base_chs', type=str, default="32,32,16", help='cost regularization base channels')
parser.add_argument('--grad_method', type=str, default="detach", choices=["detach", "undetach"], help='predicted disp detach, undetach')
from torch.utils.data import DataLoader
import gc
from PIL import Image
cudnn.benchmark = True
parser = argparse.ArgumentParser(description='seg')
parser.add_argument('--mode', type=str, default='test', help='train or test')
parser.add_argument('--model',
default='seg',
help='select a model structure',
choices=__models__.keys())
parser.add_argument('--dataset',
required=True,
help='dataset name',
choices=__datasets__.keys())
parser.add_argument('--datapath', default='', help='data path')
parser.add_argument('--channels',
type=int,
default=3,
help='net input channels')
parser.add_argument('--out_channels',
type=int,
default=1,
help='net output channels')
parser.add_argument('--testlist', required=True, help='testing list')
parser.add_argument('--test_batch_size',
type=int,
default=8,
help='testing batch size')
parser.add_argument('--test_crop_height',
from models import __models__
from utils import *
from torch.utils.data import DataLoader
import gc
from skimage import io
from matplotlib import pyplot as plt
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
cudnn.benchmark = True
parser = argparse.ArgumentParser(description='Group-wise Correlation Stereo Network (GwcNet)')
parser.add_argument('--model', default='gwcnet-gc', help='select a model structure', choices=__models__.keys())
parser.add_argument('--maxdisp', type=int, default=192, help='maximum disparity')
parser.add_argument('--dataset', default='kitti', help='dataset name', choices=__datasets__.keys())
parser.add_argument('--datapath', default='/media/data1/dh/DataSet/SceneFlowData/KITTI2015', help='data path')
parser.add_argument('--testlist', default='./filenames/kitti15_test.txt', help='testing list')
parser.add_argument('--loadckpt', default='./checkpoints/kitti15/bm_1/checkpoint_000799.ckpt', help='load the weights from a specific checkpoint')
# parse arguments
args = parser.parse_args()
# dataset, dataloader
StereoDataset = __datasets__[args.dataset]
test_dataset = StereoDataset(args.datapath, args.testlist, False)
TestImgLoader = DataLoader(test_dataset, 2, shuffle=False, num_workers=4, drop_last=False)
# # model, optimizer
model = __models__[args.model](args.maxdisp)
model = nn.DataParallel(model)