Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -20,7 +20,7 @@ transform = transforms.Compose([
|
|
20 |
def predict(image):
|
21 |
image = transform(image).unsqueeze(0)
|
22 |
with torch.no_grad():
|
23 |
-
output =
|
24 |
probabilities = F.softmax(output, dim=1)
|
25 |
class_idx = torch.argmax(probabilities, dim=1).item()
|
26 |
return "Diagram" if class_idx == 0 else "Not Diagram"
|
|
|
20 |
def predict(image):
|
21 |
image = transform(image).unsqueeze(0)
|
22 |
with torch.no_grad():
|
23 |
+
output = model(image)
|
24 |
probabilities = F.softmax(output, dim=1)
|
25 |
class_idx = torch.argmax(probabilities, dim=1).item()
|
26 |
return "Diagram" if class_idx == 0 else "Not Diagram"
|