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")