Raaniel commited on
Commit
f7f3a4a
1 Parent(s): 6d1b15c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -0
app.py CHANGED
@@ -38,6 +38,12 @@ model.load_state_dict(state)
38
  device = torch.device('cuda' if USE_CUDA else 'cpu')
39
  model = model.to(device)
40
 
 
 
 
 
 
 
41
  classes = ["chmury", 'inne', "dym"]
42
 
43
  def predict(image, model=model, classes=classes, device=device, transform=transform):
 
38
  device = torch.device('cuda' if USE_CUDA else 'cpu')
39
  model = model.to(device)
40
 
41
+ transform = torchvision.transforms.Compose([
42
+ torchvision.transforms.Resize((224, 224)),
43
+ torchvision.transforms.ToTensor(),
44
+ torchvision.transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.225, 0.225, 0.225])
45
+ ])
46
+
47
  classes = ["chmury", 'inne', "dym"]
48
 
49
  def predict(image, model=model, classes=classes, device=device, transform=transform):