parser_train = subparsers.add_parser('train')
parser_train.add_argument('--port', type=int, default=80)
parser_train.add_argument('--root', required=True)
parser_train.add_argument('--epochs', type=int, default=100)
parser_train.add_argument('--workers', type=int, default=4)
parser_train.add_argument('--batch', type=int, default=1)
parser_train.add_argument('--steps-loss', type=int, default=50)
parser_train.add_argument('--steps-plot', type=int, default=0)
parser_train.add_argument('--steps-save', type=int, default=500)
return parser.parse_args()
NUM_CHANNELS = 3
NUM_CLASSES = 2
color_transform = Colorize()
image_transform = ToPILImage()
input_transform = Compose([
Scale(512),
CenterCrop(512),
ToTensor(),
Normalize([.485, .456, .406], [.229, .224, .225]),
])
target_transform = Compose([
Scale(512),
CenterCrop(512),
ToLabel(),
Relabel(255, 1)
])
from torch.utils.data import DataLoader
from torchvision.transforms import Compose, Normalize
from torchvision.transforms import ToTensor, ToPILImage, Resize
sys.path.append("/home/nico/PycharmProjects/project-marvel/defect-detection")
from defect_detection.evaluator.evaluation import save_metrics_on_results
from piwise.criterion import CrossEntropyLoss2d
from piwise.dataset import VOCTrain, VOCTest
from piwise.network import FCN8, FCN16, FCN32, UNet, PSPNet, SegNet
from piwise.transform import ToLabel, Colorize
from piwise.visualize import Dashboard
NUM_CHANNELS = 3
NUM_CLASSES = 16
color_transform = Colorize(n=NUM_CLASSES)
image_transform = ToPILImage()
input_transform = Compose([
Resize(256),
ToTensor(),
# Normalize([.485, .456, .406], [.229, .224, .225]),
Normalize([.5, .5, .5], [.5, .5, .5]),
])
target_transform = Compose([
Resize(256),
ToLabel(),
# Relabel(255, 21),
])
def train(args, model):
