Spaces:
Running
Running
Update app.py
Browse files
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):
|