Update app.py

app.py

@@ -1,11 +1,18 @@
 import gradio as gr
-from transformers import AutoProcessor, AutoModelForImageClassification
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 import torch
 from PIL import Image
 
 model_name = 'e1010101/vit-384-tongue-image'
-processor = AutoProcessor.from_pretrained(model_name)
-model = AutoModelForImageClassification.from_pretrained(model_name)
+processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-384")
+model = AutoModelForImageClassification.from_pretrained(
+    model_name,
+    num_labels=3,
+    problem_type="multi_label_classification",
+    ignore_mismatched_sizes=True,
+    id2label={0: 'Crack', 1: 'Red-Dots', 2: 'Toothmark'},
+    label2id={'Crack': 0, 'Red-Dots': 1, 'Toothmark': 2}
+)
 
 def classify_image(image):
     inputs = processor(images=image, return_tensors="pt")