czkaiweb commited on
Commit
7f09978
1 Parent(s): 05e6243

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -5
app.py CHANGED
@@ -12,19 +12,19 @@ os.system("wget https://github.com/liuxiaoyuyuyu/vanGogh-and-Other-Artist/blob/m
12
  #model = torch.hub.load('pytorch/vision:v0.9.0', 'mobilenet_v2', pretrained=False)
13
  #checkpoint = 'https://github.com/liuxiaoyuyuyu/vanGogh-and-Other-Artist/blob/main/model_weights_mobilenet_v2_valp1trainp2.pth'
14
  #model.load_state_dict(torch.hub.load_state_dict_from_url(checkpoint, progress=False))
15
- model = models.mobilenet_v2()
16
- num_ftrs = model.classifier[1].in_features
17
- model.classifier[1] = nn.Linear(num_ftrs, 6)
18
  device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
19
  #model = model.to(device)
20
- model.load_state_dict(torch.load('model_weights_mobilenet_v2_valp1trainp2.pth',map_location=device))
21
 
22
  #torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
23
 
24
 
25
  def inference(input_image):
26
  preprocess = transforms.Compose([
27
- transforms.Resize(256),
28
  transforms.CenterCrop(224),
29
  transforms.ToTensor(),
30
  #transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 
12
  #model = torch.hub.load('pytorch/vision:v0.9.0', 'mobilenet_v2', pretrained=False)
13
  #checkpoint = 'https://github.com/liuxiaoyuyuyu/vanGogh-and-Other-Artist/blob/main/model_weights_mobilenet_v2_valp1trainp2.pth'
14
  #model.load_state_dict(torch.hub.load_state_dict_from_url(checkpoint, progress=False))
15
+ model = models.vgg16()
16
+ num_ftrs = model.classifier[6].in_features
17
+ model.classifier[6] = nn.Linear(num_ftrs, 6)
18
  device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
19
  #model = model.to(device)
20
+ model.load_state_dict(torch.load('VGG16_weights_May28.pth',map_location=device))
21
 
22
  #torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
23
 
24
 
25
  def inference(input_image):
26
  preprocess = transforms.Compose([
27
+ transforms.Resize(260),
28
  transforms.CenterCrop(224),
29
  transforms.ToTensor(),
30
  #transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),