Update app.py

app.py

@@ -5,6 +5,10 @@ from torchvision.transforms import Compose, Resize, ToTensor, Normalize, RandomH
 from PIL import Image
 import traceback
 
+# Load dataset to get labels
+dataset = load_dataset("bentrevett/caltech-ucsd-birds-200-2011")
+labels = dataset['train'].features['label'].names
+
 # Load model and processor
 model_name = "riyadifirman/klasifikasiburung"
 processor = AutoImageProcessor.from_pretrained(model_name)
@@ -27,7 +31,8 @@ def predict(image):
         outputs = model(inputs)
         logits = outputs.logits
         predicted_class_idx = logits.argmax(-1).item()
-        return processor.decode(predicted_class_idx)
+        predicted_class = labels[predicted_class_idx]
+        return predicted_class
     except Exception as e:
         # Menampilkan error
         print("An error occurred:", e)
@@ -44,4 +49,4 @@ interface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch(share=True)