Update app.py

app.py

@@ -1,6 +1,7 @@
 import torch
 from torchvision import transforms
 import gradio as gr
+from PIL import Image, ImageDraw
 
 # Carregue o dicionário contendo o modelo PyTorch treinado
 model_dict = torch.load("best.pt", map_location=torch.device('cpu'))  # Use 'cpu' se não estiver usando GPU
@@ -13,7 +14,6 @@ model.eval()
 
 # Transformação de pré-processamento
 preprocess = transforms.Compose([
-    transforms.ToPILImage(),
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
@@ -29,12 +29,34 @@ def predict(image):
     with torch.no_grad():
         output = model(input_batch)
 
-    # Post-processamento, se necessário
-    # ...
+    # Post-processamento para extrair informações de detecção
+    # (substitua as linhas abaixo de acordo com a estrutura de saída do seu modelo)
+    boxes = output['boxes']
+    labels = output['labels']
+    scores = output['scores']
 
-    # Retorna o resultado da inferência
-    return output
+    # Visualize a imagem com caixas de detecção
+    image_with_boxes = visualize_detections(image, boxes, labels, scores)
+
+    # Retorna a imagem com caixas de detecção
+    return image_with_boxes
+
+# Função para visualizar as caixas de detecção na imagem
+def visualize_detections(image, boxes, labels, scores):
+    # Converta a imagem para o formato PIL
+    image_pil = transforms.ToPILImage()(image)
+
+    # Crie um objeto ImageDraw para desenhar caixas na imagem
+    draw = ImageDraw.Draw(image_pil)
+
+    # Desenhe as caixas de detecção na imagem
+    for box, label, score in zip(boxes, labels, scores):
+        box = [round(coord, 2) for coord in box.tolist()]  # Arredonde as coordenadas da caixa
+        draw.rectangle(box, outline="red", width=3)
+        draw.text((box[0], box[1]), f"Label: {label}\nScore: {score:.2f}", fill="red")
+
+    return image_pil
 
 # Interface Gradio
-iface = gr.Interface(fn=predict, inputs="image", outputs="text")
+iface = gr.Interface(fn=predict, inputs="image", outputs="image")
 iface.launch()