Update app.py

app.py

@@ -5,18 +5,22 @@ import matplotlib.pyplot as plt
 import seaborn as sns
 import cv2
 import torch
-from transformers import AutoTokenizer, AutoModelForImageClassification
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, login
 from PIL import Image
 import pickle
 
+
 # Cargar dataset
 dataset_repo = "VicGerardoPR/Traffic_sign_dataset"
 dataset_file = "train.p"
 
-dataset_path = hf_hub_download(repo_id=dataset_repo, filename=dataset_file)
-with open(dataset_path, 'rb') as file:
-    train_data = pickle.load(file)
+try:
+    dataset_path = hf_hub_download(repo_id=dataset_repo, filename=dataset_file, repo_type="dataset")
+    with open(dataset_path, 'rb') as file:
+        train_data = pickle.load(file)
+except Exception as e:
+    st.error(f"Error al cargar el dataset: {e}")
+    st.stop()
 
 train_images, train_labels = train_data['features'], train_data['labels']
 
@@ -32,9 +36,14 @@ train_images, train_labels = preprocess_data(train_images, train_labels)
 model_repo = "VicGerardoPR/Traffic_sign_model"
 model_file = "traffic_sign_classifier.h5"
 
-model_path = hf_hub_download(repo_id=model_repo, filename=model_file)
-model = torch.load(model_path, map_location=torch.device('cpu'))
-model.eval()
+try:
+    model_path = hf_hub_download(repo_id=model_repo, filename=model_file, repo_type="model")
+    model = torch.load(model_path, map_location=torch.device('cpu'))
+    model.eval()
+except Exception as e:
+    st.error(f"Error al cargar el modelo: {e}")
+    st.stop()
+
 
 # Diccionario de clases de señales de tráfico
 classes = {