Update app.py

app.py

@@ -5,7 +5,7 @@ import pandas as pd
 import matplotlib.pyplot as plt
 import seaborn as sns
 import cv2
-import torch
+import tensorflow as tf
 from huggingface_hub import hf_hub_download
 from PIL import Image
 import pickle
@@ -49,13 +49,11 @@ model_file = "traffic_sign_classifier.h5"
 
 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()
+    model = tf.keras.models.load_model(model_path)
 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 = {
     0: 'Speed limit (20km/h)',
@@ -108,12 +106,9 @@ def predict(image):
     image = cv2.resize(image, (32, 32))
     image = image / 255.0
     image = np.expand_dims(image, axis=0)
-    image = np.transpose(image, (0, 3, 1, 2))  # Reordenar dimensiones para PyTorch
-    image = torch.tensor(image).float()
 
-    with torch.no_grad():
-        predictions = model(image)
-        class_idx = predictions.argmax().item()
+    predictions = model.predict(image)
+    class_idx = np.argmax(predictions, axis=1)[0]
     return classes[class_idx]
 
 # Título y descripción de la aplicación