Update app.py

app.py

@@ -29,25 +29,36 @@ def classify_image(input_image):
     outputs = model(**inputs)
     # Obter a classe prevista
     predicted_class_id = torch.argmax(outputs.logits, dim=1).item()
+    # Obter a probabilidade da classe prevista
+    predicted_class_prob = torch.softmax(outputs.logits, dim=1)[0, predicted_class_id].item()
     # Converter o ID da classe em rótulo usando o mapeamento id2label
     predicted_class_label = id2label.get(str(predicted_class_id), "Desconhecido")
     
     # Converter a imagem de numpy para BGR (formato OpenCV)
     input_image_bgr = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
     
+    # Definir cores de borda para cada classe (aqui, cores aleatórias)
+    class_colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0),
+                    (255, 0, 255), (0, 255, 255), (128, 128, 128), (0, 0, 0)]
+    
+    # Adicionar uma borda colorida à imagem
+    border_color = class_colors[predicted_class_id]
+    input_image_bgr = cv2.copyMakeBorder(input_image_bgr, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=border_color)
+    
     # Adicionar o rótulo da previsão na imagem
     font = cv2.FONT_HERSHEY_SIMPLEX
-    text = f'Previsao: {predicted_class_label}'
+    text = f'Classe: {predicted_class_label} ({predicted_class_prob:.2f})'
     text_size = cv2.getTextSize(text, font, 0.7, 2)[0]
     text_x = (input_image_bgr.shape[1] - text_size[0]) // 2
     text_y = input_image_bgr.shape[0] - 20
-    cv2.putText(input_image_bgr, text, (text_x, text_y), font, 0.7, (255, 255, 255), 2)
+    cv2.putText(input_image_bgr, text, (text_x, text_y), font, 0.7, (255, 255, 255), 2, cv2.LINE_AA)
     
     # Converter a imagem resultante de volta para RGB (formato Pillow)
     result_image = cv2.cvtColor(input_image_bgr, cv2.COLOR_BGR2RGB)
     return result_image
 
 
+
 # Criar uma interface Gradio
 interface = gr.Interface(
     fn=classify_image,