Update app.py

app.py

@@ -5,28 +5,48 @@ from torchvision import transforms, models
 import pickle
 from resnest.torch import resnest50
 
-# Carregar nomes das classes e criar mapeamento para português
+# Carregar nomes das classes originais
 with open('class_names.pkl', 'rb') as f:
     class_names_en = pickle.load(f)
 
-# Dicionário de tradução das classes para português
+# Imprimir as classes originais para debug
+print("Classes originais encontradas:", class_names_en)
+
+# Dicionário de tradução mais completo (incluindo variações)
 class_names_pt = {
     'apple': 'maçã',
+    'Apple': 'maçã',
+    'Apple 10': 'maçã',  # adicionando variações
     'banana': 'banana',
+    'Banana': 'banana',
     'cherry': 'cereja',
-    'chico': 'sapoti',  # nome em português para chico/fruit
+    'Cherry': 'cereja',
+    'chico': 'sapoti',
     'grape': 'uva',
+    'Grape': 'uva',
     'kiwi': 'kiwi',
+    'Kiwi': 'kiwi',
     'mango': 'manga',
+    'Mango': 'manga',
     'orange': 'laranja',
+    'Orange': 'laranja',
     'pear': 'pera',
-    'tomato': 'tomate'
+    'Pear': 'pera',
+    'tomato': 'tomate',
+    'Tomato': 'tomate'
 }
 
-# Criar lista de nomes em português na mesma ordem que class_names_en
-class_names = [class_names_pt[en] for en in class_names_en]
+# Criar lista de nomes em português, usando o nome original se não houver tradução
+class_names = []
+for en in class_names_en:
+    # Remover números e espaços extras para normalizar
+    base_name = ''.join([i for i in en if not i.isdigit()]).strip()
+    translated = class_names_pt.get(base_name, class_names_pt.get(en, en))
+    class_names.append(translated)
+
+print("Classes traduzidas:", class_names)
 
-# Carregar o modelo treinado
+# Restante do código permanece igual...
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 model = resnest50(pretrained=None)
@@ -68,13 +88,13 @@ def predict_image(img):
     top3_probs, top3_indices = torch.topk(probabilities, 3)
 
     results = {
-        class_names[i]: p.item()
+        class_names[i]: float(p)
         for p, i in zip(top3_probs, top3_indices)
     }
 
     # Obter a melhor previsão
     best_class = class_names[top3_indices[0]]
-    best_conf = top3_probs[0].item() * 100
+    best_conf = float(top3_probs[0]) * 100
 
     # Salvar resultados
     with open('/tmp/prediction_results.txt', 'a') as f: