Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -1,14 +1,14 @@
|
|
1 |
import gradio as gr
|
2 |
from PIL import Image
|
3 |
from transformers import ViTImageProcessor, ViTForImageClassification
|
4 |
-
|
5 |
|
6 |
|
7 |
processor = ViTImageProcessor.from_pretrained('Rageshhf/fine-tuned-model')
|
8 |
|
9 |
id2label = {0: 'Mild_Demented', 1: 'Moderate_Demented', 2: 'Non_Demented', 3: 'Very_Mild_Demented'}
|
10 |
label2id = {'Mild_Demented': 0, 'Moderate_Demented': 1, 'Non_Demented': 2, 'Very_Mild_Demented': 3}
|
11 |
-
|
12 |
|
13 |
model = ViTForImageClassification.from_pretrained(
|
14 |
'Rageshhf/fine-tuned-model',
|
@@ -23,15 +23,20 @@ description = """Trained to classify disease based on image data."""
|
|
23 |
|
24 |
|
25 |
def predict(image):
|
26 |
-
|
27 |
inputs = processor(images=image, return_tensors="pt")
|
28 |
outputs = model(**inputs)
|
|
|
29 |
logits = outputs.logits
|
30 |
-
|
31 |
-
|
32 |
-
|
|
|
|
|
|
|
|
|
|
|
33 |
|
34 |
-
demo = gr.Interface(fn=predict, inputs="image", outputs=gr.Label(num_top_classes=
|
35 |
description=description,).launch()
|
36 |
|
37 |
# demo.launch(debug=True)
|
|
|
1 |
import gradio as gr
|
2 |
from PIL import Image
|
3 |
from transformers import ViTImageProcessor, ViTForImageClassification
|
4 |
+
import torch
|
5 |
|
6 |
|
7 |
processor = ViTImageProcessor.from_pretrained('Rageshhf/fine-tuned-model')
|
8 |
|
9 |
id2label = {0: 'Mild_Demented', 1: 'Moderate_Demented', 2: 'Non_Demented', 3: 'Very_Mild_Demented'}
|
10 |
label2id = {'Mild_Demented': 0, 'Moderate_Demented': 1, 'Non_Demented': 2, 'Very_Mild_Demented': 3}
|
11 |
+
labels = ['Mild_Demented', 'Moderate_Demented', 'Non_Demented', 'Very_Mild_Demented']
|
12 |
|
13 |
model = ViTForImageClassification.from_pretrained(
|
14 |
'Rageshhf/fine-tuned-model',
|
|
|
23 |
|
24 |
|
25 |
def predict(image):
|
|
|
26 |
inputs = processor(images=image, return_tensors="pt")
|
27 |
outputs = model(**inputs)
|
28 |
+
|
29 |
logits = outputs.logits
|
30 |
+
prediction = torch.nn.functional.softmax(logits, dim=1)
|
31 |
+
probabilities = prediction[0].tolist()
|
32 |
+
|
33 |
+
output = {}
|
34 |
+
for i, prob in enumerate(probabilities):
|
35 |
+
output[labels[i]] = prob
|
36 |
+
|
37 |
+
return output
|
38 |
|
39 |
+
demo = gr.Interface(fn=predict, inputs="image", outputs=gr.Label(num_top_classes=3), title=title, examples=["examples/image_1.png", "examples/image_2.png", "examples/image_3.png"],
|
40 |
description=description,).launch()
|
41 |
|
42 |
# demo.launch(debug=True)
|