Update app.py
Browse files
app.py
CHANGED
@@ -27,7 +27,7 @@ if file_name is not None:
|
|
27 |
# Convert grayscale image to RGB format
|
28 |
image_tensor = extractor(images=image, return_tensors="pt")["pixel_values"]
|
29 |
predictions = model(image_tensor)
|
30 |
-
predicted_class = torch.argmax(
|
31 |
col2.header("Probabilities")
|
32 |
col2.subheader(predicted_class.item())
|
33 |
|
|
|
27 |
# Convert grayscale image to RGB format
|
28 |
image_tensor = extractor(images=image, return_tensors="pt")["pixel_values"]
|
29 |
predictions = model(image_tensor)
|
30 |
+
predicted_class = torch.argmax(predictions.logits, dim=1)
|
31 |
col2.header("Probabilities")
|
32 |
col2.subheader(predicted_class.item())
|
33 |
|