Update app.py

app.py

@@ -5,10 +5,28 @@ from torchvision import transforms, models
 import pickle
 from resnest.torch import resnest50
 
+# Carregar nomes das classes e criar mapeamento para português
 with open('class_names.pkl', 'rb') as f:
-    class_names = pickle.load(f)
-
-# 加载训练好的模型
+    class_names_en = pickle.load(f)
+
+# Dicionário de tradução das classes para português
+class_names_pt = {
+    'apple': 'maçã',
+    'banana': 'banana',
+    'cherry': 'cereja',
+    'chico': 'sapoti',  # nome em português para chico/fruit
+    'grape': 'uva',
+    'kiwi': 'kiwi',
+    'mango': 'manga',
+    'orange': 'laranja',
+    'pear': 'pera',
+    '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]
+
+# Carregar o modelo treinado
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 model = resnest50(pretrained=None)
@@ -17,12 +35,12 @@ model.fc = nn.Sequential(
     nn.Linear(model.fc.in_features, len(class_names))
 )
 
-# 加载模型权重
+# Carregar os pesos do modelo
 model.load_state_dict(torch.load('best_model.pth', map_location=device, weights_only=True))
 model = model.to(device)
 model.eval()
 
-# 定义与训练时相同的预处理流程
+# Definir o mesmo pré-processamento usado no treinamento
 preprocess = transforms.Compose([
     transforms.Resize((100, 100)),
     transforms.ToTensor(),
@@ -33,20 +51,20 @@ preprocess = transforms.Compose([
 def predict_image(img):
     img = img.convert('RGB')
 
-    # 应用预处理
+    # Aplicar pré-processamento
     input_tensor = preprocess(img)
 
-    # 添加批次维度并移动到设备
+    # Adicionar dimensão de batch e mover para o dispositivo
     input_batch = input_tensor.unsqueeze(0).to(device)
 
-    # 预测
+    # Fazer previsão
     with torch.no_grad():
         output = model(input_batch)
 
-    # 计算概率
+    # Calcular probabilidades
     probabilities = torch.nn.functional.softmax(output[0], dim=0)
 
-    # 获取前3个预测结果
+    # Obter as 3 melhores previsões
     top3_probs, top3_indices = torch.topk(probabilities, 3)
 
     results = {
@@ -54,42 +72,42 @@ def predict_image(img):
         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
 
-    # 保存结果
+    # Salvar resultados
     with open('/tmp/prediction_results.txt', 'a') as f:
-        f.write(f"Image: {img}\n"
-                f"Predicted: {best_class}\n"
-                f"Confidence: {best_conf:.2f}%\n"
+        f.write(f"Imagem: {img}\n"
+                f"Previsão: {best_class}\n"
+                f"Confiança: {best_conf:.2f}%\n"
                 f"Top 3: {results}\n"
                 f"------------------------\n")
 
-    return best_class, best_conf, results
+    return best_class, f"{best_conf:.2f}%", results
 
-# 创建Gradio界面
+# Criar interface Gradio
 def create_interface():
     examples = [
         "r0_0_100.jpg",
         "r0_18_100.jpg"
     ]
 
-    with gr.Blocks(title="Fruit Classification", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 🍎 水果识别系统")
+    with gr.Blocks(title="Sistema de Classificação de Frutas", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# 🍎 Sistema de Reconhecimento de Frutas")
 
         with gr.Row():
             with gr.Column():
-                image_input = gr.Image(type="pil", label="上传图像")
+                image_input = gr.Image(type="pil", label="Envie uma imagem")
                 gr.Examples(examples=examples, inputs=image_input)
-                submit_btn = gr.Button("分类", variant="primary")
+                submit_btn = gr.Button("Classificar", variant="primary")
 
             with gr.Column():
-                best_pred = gr.Textbox(label="预测结果")
-                confidence = gr.Textbox(label="置信度")
+                best_pred = gr.Textbox(label="Resultado da Previsão")
+                confidence = gr.Textbox(label="Nível de Confiança")
                 full_results = gr.Label(label="Top 3", num_top_classes=3)
 
-        # ‘分类’按钮点击事件
+        # Evento de clique do botão 'Classificar'
         submit_btn.click(
             fn=predict_image,
             inputs=image_input,
@@ -101,4 +119,4 @@ def create_interface():
 
 if __name__ == "__main__":
     interface = create_interface()
-    interface.launch(share=False)
+    interface.launch(share=False)