def __init__(self, args):
self.args = args
# image transform
input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([.485, .456, .406], [.229, .224, .225]),
])
# dataset and dataloader
data_kwargs = {'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size}
train_dataset = get_segmentation_dataset(args.dataset, split=args.train_split, mode='train', **data_kwargs)
val_dataset = get_segmentation_dataset(args.dataset, split='val', mode='val', **data_kwargs)
self.train_loader = data.DataLoader(dataset=train_dataset,
batch_size=args.batch_size,
shuffle=True,
drop_last=True)
self.val_loader = data.DataLoader(dataset=val_dataset,
batch_size=1,
shuffle=False)
# create network
self.model = get_fast_scnn(dataset=args.dataset, aux=args.aux)
if torch.cuda.device_count() > 1:
self.model = torch.nn.DataParallel(self.model, device_ids=[0, 1, 2])
self.model.to(args.device)
# resume checkpoint if needed
if args.resume:
if os.path.isfile(args.resume):
name, ext = os.path.splitext(args.resume)
assert ext == '.pkl' or '.pth', 'Sorry only .pth and .pkl files supported.'
print('Resuming training, loading {}...'.format(args.resume))
self.model.load_state_dict(torch.load(args.resume, map_location=lambda storage, loc: storage))
# create criterion
self.criterion = MixSoftmaxCrossEntropyOHEMLoss(aux=args.aux, aux_weight=args.aux_weight,
ignore_index=-1).to(args.device)
# optimizer
self.optimizer = torch.optim.SGD(self.model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# lr scheduling
self.lr_scheduler = LRScheduler(mode='poly', base_lr=args.lr, nepochs=args.epochs,
iters_per_epoch=len(self.train_loader),
power=0.9)
# evaluation metrics
self.metric = SegmentationMetric(train_dataset.num_class)
self.best_pred = 0.0
def __init__(self, args):
self.args = args
# output folder
self.outdir = 'test_result'
if not os.path.exists(self.outdir):
os.makedirs(self.outdir)
# image transform
input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([.485, .456, .406], [.229, .224, .225]),
])
# dataset and dataloader
val_dataset = get_segmentation_dataset(args.dataset,
split='val',
mode='test',
transform=input_transform)
self.val_loader = data.DataLoader(dataset=val_dataset,
batch_size=1,
shuffle=False)
# create network
self.model = get_fast_scnn(args.dataset,
aux=args.aux,
pretrained=True,
root=args.save_folder).to(args.device)
print('Finished loading model!')
self.metric = SegmentationMetric(val_dataset.num_class)
def __init__(self, args):
self.args = args
# output folder
self.outdir = 'test_result'
if not os.path.exists(self.outdir):
os.makedirs(self.outdir)
# image transform
input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([.485, .456, .406], [.229, .224, .225]),
])
# dataset and dataloader
val_dataset = get_segmentation_dataset(args.dataset,
split='val',
mode='testval',
transform=input_transform)
self.val_loader = data.DataLoader(dataset=val_dataset,
batch_size=1,
shuffle=False)
# create network
self.model = get_model().to(self.args.device)
self.model.load_state_dict(
torch.load(
"/home/rajiv/Documents/papers/API-Semantic-Segmentation/weights/psp_citys.pth"
))
print('Finished loading model!')
self.metric = SegmentationMetric(val_dataset.num_class)
def __init__(self,NUM_CHANNELS):
MODEL_NAME = "fast_scnn" if NUM_CHANNELS==4 else "fast_scnn_rgb"
self.model_name = MODEL_NAME
self.outdir = ARTIFACT_DETECTION_DIR + "/inference/" + self.model_name
if not os.path.exists(self.outdir):
os.makedirs(self.outdir)
if MODEL_NAME == "fast_scnn":
input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([.485, .456, .406,0.4], [.229, .224, .225,0.4]),
])
else:
input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([.485, .456, .406], [.229, .224, .225]),
])
val_dataset = get_segmentation_dataset(num_channels=NUM_CHANNELS,root=ARGS_INFERENCE_DIR, split='val', mode='testval', transform=input_transform)
self.val_loader = data.DataLoader(dataset=val_dataset,
batch_size=1,
shuffle=False)
WEIGHTS_PATH = ARTIFACT_DETECTION_DIR+"/weights/fast_scnn_rgb" if NUM_CHANNELS==3 else ARTIFACT_DETECTION_DIR+"/weights/fast_scnn"
print("Weights Path:", WEIGHTS_PATH)
self.model = get_fast_scnn("fast_scnn",num_channels=NUM_CHANNELS, aux=False, pretrained=True, root=WEIGHTS_PATH).cuda()
self.metric = SegmentationMetric(val_dataset.num_class)
def __init__(self, args, cfg=None):
# train_dataset = [build_dataset(cfg.data.train)]
# self.dataset= train_dataset
# val_dataset = [build_dataset(cfg.data.test)]
# if len(cfg.workflow) == 2:
# train_dataset.append(build_dataset(cfg.data.val))
# train_data_loaders = [
# build_dataloader(
# ds,
# cfg.data.imgs_per_gpu,
# cfg.data.workers_per_gpu,
# # cfg.gpus,
# dist=False) for ds in train_dataset
# ]
# val_data_loader = build_dataloader(
# val_dataset,
# imgs_per_gpu=1,
# workers_per_gpu=cfg.data.workers_per_gpu,
# dist=False,
# shuffle=False)
# self.train_loader = train_data_loaders[0]
# self.val_loader = val_data_loader
self.args = args
# image transform
input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[123.675, 116.28, 103.53],
std=[58.395, 57.12, 57.375]),
])
# dataset and dataloader
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size
}
train_dataset = get_segmentation_dataset(args.dataset,
split=args.train_split,
mode='train',
**data_kwargs)
val_dataset = get_segmentation_dataset(args.dataset,
split='val',
mode='val',
**data_kwargs)
self.train_loader = data.DataLoader(dataset=train_dataset,
batch_size=args.batch_size,
shuffle=True,
drop_last=True)
self.val_loader = data.DataLoader(dataset=val_dataset,
batch_size=1,
shuffle=False)
# create network
self.model = get_fast_scnn(dataset=args.dataset, aux=args.aux)
if torch.cuda.device_count() > 1:
self.model = torch.nn.DataParallel(self.model,
device_ids=[0, 1, 2])
self.model.to(args.device)
# resume checkpoint if needed
if args.resume:
if os.path.isfile(args.resume):
name, ext = os.path.splitext(args.resume)
assert ext == '.pkl' or '.pth', 'Sorry only .pth and .pkl files supported.'
print('Resuming training, loading {}...'.format(args.resume))
self.model.load_state_dict(
torch.load(args.resume,
map_location=lambda storage, loc: storage))
# create criterion
self.criterion = MixSoftmaxCrossEntropyOHEMLoss(
aux=args.aux, aux_weight=args.aux_weight,
ignore_index=-1).to(args.device)
# optimizer
self.optimizer = torch.optim.SGD(self.model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# lr scheduling
self.lr_scheduler = LRScheduler(mode='poly',
base_lr=args.lr,
nepochs=args.epochs,
iters_per_epoch=len(self.train_loader),
power=0.9)
# evaluation metrics
self.metric = SegmentationMetric(train_dataset.num_class)
self.best_pred = 0.0
from torchvision import transforms
from data_loader import get_segmentation_dataset
from models.fast_scnn import get_fast_scnn
from utils.loss import MixSoftmaxCrossEntropyLoss, MixSoftmaxCrossEntropyOHEMLoss
from utils.lr_scheduler import LRScheduler
from utils.metric import SegmentationMetric
from mmcv import Config
from data_loader.datasets import build_dataset, build_dataloader
from tools_zym.weight import bbox2multimask
import os
import PIL
import numpy as np
train_dataset = get_segmentation_dataset('coco', split='train', mode='train')
val_dataset = get_segmentation_dataset('coco', split='val', mode='val')
root = os.path.join('./data/mask')
if not os.path.exists(root):
os.mkdir(root)
os.mkdir(os.path.join(root, 'train2017'))
os.mkdir(os.path.join(root, 'val2017'))
# for i, (img, mask) in enumerate(train_dataset):
# print(mask)
for i, (img, mask) in enumerate(train_dataset):
ids = train_dataset.img_ids[i]
image_info = train_dataset.coco.loadImgs(ids)[0]