forked from simo-an/FastFlow
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
29 lines (23 loc) · 997 Bytes
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
'''This is the repo which contains the original code to the WACV 2021 paper
"Same Same But DifferNet: Semi-Supervised Defect Detection with Normalizing Flows"
by Marco Rudolph, Bastian Wandt and Bodo Rosenhahn.
For further information contact Marco Rudolph (rudolph@tnt.uni-hannover.de)'''
import config as c
from train import train
from utils import load_datasets, make_dataloaders
import os
import sys
print("Python version")
print (sys.version)
print(os.listdir())
print(os.listdir("./data"))
#I change the location where pytorch saves pretrained models
os.environ['TORCH_HOME'] = 'models\\alexnet' #setting the environment variable
#import torch
#xcv = torch.hub.load('facebookresearch/deit:main', 'deit_base_distilled_patch16_224', pretrained=True)
#manage dataset
from handledata import handledata
#handledata()
train_set, test_set = load_datasets(c.dataset_path, c.class_name)
train_loader, test_loader = make_dataloaders(train_set, test_set)
model = train(train_loader, test_loader)