go3

app.py

@@ -3,43 +3,67 @@ import numpy as np
 from PIL import Image
 import torch
 import matplotlib.pyplot as plt
+import io
 
 # Modelo
-model = torch.hub.load('ultralytics/yolov5', 'custom', path='bestyolo5.pt')
+model = torch.hub.load('ultralytics/yolov5', 'custom', path='bestyolo5.pt', trust_repo=True)
 
 def detect(img):
+    if img is None:
+        return None, 0
+    
     img_arr = np.array(img)
     results = model(img_arr)
-
-    fig, ax = plt.subplots()
+    
+    # Criar figura
+    fig, ax = plt.subplots(figsize=(10, 8))
     ax.imshow(img_arr)
-
+    
     cattle_count = 0
+    
+    # Processar detecções
     for *xyxy, conf, cls in results.xyxy[0].cpu().numpy():
         x1, y1, x2, y2 = map(int, xyxy)
         label = model.names[int(cls)]
+        
         if label == 'cattle':
             cattle_count += 1
-        ax.add_patch(plt.Rectangle((x1, y1), x2-x1, y2-y1, fill=False, color='red', linewidth=2))
-        ax.text(x1, y1, f'{label} {conf:.2f}', color='white', fontsize=8, bbox={'facecolor': 'red', 'alpha': 0.5})
-
+            
+        # Desenhar bounding box
+        ax.add_patch(plt.Rectangle((x1, y1), x2-x1, y2-y1, 
+                                 fill=False, color='red', linewidth=2))
+        
+        # Adicionar label
+        ax.text(x1, y1-10, f'{label} {conf:.2f}', 
+               color='white', fontsize=10, 
+               bbox={'facecolor': 'red', 'alpha': 0.7})
+    
+    ax.set_title(f'Detecções: {cattle_count} gado encontrado')
     plt.axis('off')
-
-    fig.canvas.draw()
-    pil_img = Image.fromarray(np.array(fig.canvas.renderer._renderer))
+    plt.tight_layout()
+    
+    # Converter para PIL Image
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', bbox_inches='tight', dpi=150)
+    buf.seek(0)
+    pil_img = Image.open(buf)
     plt.close(fig)
-
+    
     return pil_img, cattle_count
 
-# Gradio
+# Interface Gradio
 iface = gr.Interface(
     fn=detect,
     inputs=gr.Image(type="pil"),
-    outputs=[gr.Image(type="pil"), gr.Textbox(label="Number of Cattle Detected")],
-    title="YOLOv5 Cattle Counter",
-    description="Object detector trained to count cattle using YOLOv5.",
-    examples=[["example1.jpg"]]
+    outputs=[
+        gr.Image(type="pil", label="Imagem com Detecções"), 
+        gr.Number(label="Número de Gado Detectado")
+    ],
+    title="🐄 YOLOv5 Contador de Gado",
+    description="Detector de objetos treinado para contar gado usando YOLOv5.",
+    examples=[["example1.jpg"]] if False else None,  # Remova ou adicione imagens de exemplo
+    theme=gr.themes.Soft()
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch(share=True, debug=True)

requirements.txt

@@ -2,6 +2,11 @@ gradio==4.29.0
 numpy
 Pillow
 torch
+torchvision
 matplotlib
 opencv-python
 seaborn
+ultralytics
+PyYAML
+requests
+tqdm