Shchushch commited on
Commit
7ad97f7
1 Parent(s): 40f8d43

Update assets/skin.py

Browse files
Files changed (1) hide show
  1. assets/skin.py +5 -12
assets/skin.py CHANGED
@@ -2,14 +2,8 @@
2
  import torch
3
  from torchvision import transforms
4
  from torchvision.models import VGG19_BN_Weights, vgg19_bn
5
-
6
- model =torch.load('assets/skin.pth')
7
- #print(torch.load('assets/weights.pth'))
8
- # model=vgg19_bn()
9
- # print(torch.load('assets/model.pth'))
10
- # model.load_state_dict(torch.load('assets/skin_cancer.pth'))
11
- # print(model)
12
- # torch.save(model.state_dict(), 'assets/skin_cancer.pth')
13
  skin_map={0:'Добро',1:'Зло'}
14
  def get_evil(img):
15
 
@@ -20,12 +14,11 @@ def get_evil(img):
20
 
21
  input_image = transform(img).unsqueeze(0) # Добавьте размерность пакета (batch dimension)
22
 
23
- device = 'cpu'
24
  #torch.device("cuda" if torch.cuda.is_available() else 'cpu')
25
- model.to(device)
26
  model.eval()
27
- input_image = input_image.to(device)
28
- #model.to(device)
29
  with torch.no_grad():
30
  res=model(input_image).item()
31
  return f'Степень злобы: {res}\n\n Т.е. опухоль: {skin_map[round(res)]}'
 
2
  import torch
3
  from torchvision import transforms
4
  from torchvision.models import VGG19_BN_Weights, vgg19_bn
5
+ DEVICE = 'cpu'
6
+ model =torch.load('assets/skin.pth', map_location=)
 
 
 
 
 
 
7
  skin_map={0:'Добро',1:'Зло'}
8
  def get_evil(img):
9
 
 
14
 
15
  input_image = transform(img).unsqueeze(0) # Добавьте размерность пакета (batch dimension)
16
 
17
+
18
  #torch.device("cuda" if torch.cuda.is_available() else 'cpu')
19
+ model.to(DEVICE)
20
  model.eval()
21
+ input_image = input_image.to(DEVICE)
 
22
  with torch.no_grad():
23
  res=model(input_image).item()
24
  return f'Степень злобы: {res}\n\n Т.е. опухоль: {skin_map[round(res)]}'