Update app.py

app.py

@@ -1,40 +1,52 @@
-import gradio as gr
-import torch
-from torchvision import transforms
-import numpy as np
-
-# Load YOLOv7 model
-weights_path = "cattle.pt"
-config_path = "yolov7.yaml"
-net = None  # Initialize your YOLOv7 model here
-
-# Define function to detect objects using YOLOv7
-def detect_objects(image):
-    global net
-
-    if net is None:
-        # Initialize your YOLOv7 model if not already initialized
-        net = ...
-
-    # Perform any necessary preprocessing on the image
-    transform = transforms.Compose([
-        transforms.Resize((416, 416)),  # Resize image to expected input size
-        transforms.ToTensor(),           # Convert image to tensor
-    ])
-    image_tensor = transform(image).unsqueeze(0)
-
-    # Perform object detection
-    with torch.no_grad():
-        # Forward pass through the network
-        outs = net(image_tensor)
-
-        # Process detection results and draw bounding boxes (similar to your previous code)
-        ...
-
-    return image, object_count
-
-# Create Gradio interface
-inputs = gr.inputs.Image(label="Upload Image or Video")
-outputs = [gr.outputs.Image(label="Output Image with Objects Detected"), gr.outputs.Text(label="Object Count")]
-
+import gradio as gr 
+import torch 
+from torchvision import transforms 
+from models.experimental import attempt_load
+from utils.general import non_max_suppression, plot_one_box
+import numpy as np 
+
+# Load YOLOv7 model 
+weights_path = "cattle.pt" 
+config_path = "yolov7.yaml" 
+
+# Initialize your YOLOv7 model here 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+net = attempt_load(weights_path, map_location=device)
+
+# Define function to detect objects using YOLOv7 
+def detect_objects(image): 
+    global net 
+
+    # Perform any necessary preprocessing on the image 
+    transform = transforms.Compose([ 
+        transforms.Resize((416, 416)), # Resize image to expected input size 
+        transforms.ToTensor(), # Convert image to tensor 
+    ]) 
+    image_tensor = transform(image).unsqueeze(0) 
+
+    # Perform object detection 
+    with torch.no_grad(): 
+        # Forward pass through the network 
+        outs = net(image_tensor) 
+        
+        # Apply non-maximum suppression
+        pred = non_max_suppression(outs, conf_thres=0.4, iou_thres=0.5, classes=None, agnostic=False)
+        
+        # Process detection results and draw bounding boxes
+        object_count = 0
+        for i, det in enumerate(pred):
+            # Skip if no detections
+            if len(det):
+                # Increment object count
+                object_count += len(det)
+                # Loop over the detections
+                for *xyxy, conf, cls in reversed(det):
+                    # Draw bounding box
+                    plot_one_box(xyxy, image, label=f'{conf:.2f}', color=(255, 0, 0), line_thickness=3)
+                    
+    return image, object_count 
+
+# Create Gradio interface 
+inputs = gr.inputs.Image(label="Upload Image or Video") 
+outputs = [gr.outputs.Image(label="Output Image with Objects Detected"), gr.outputs.Text(label="Object Count")] 
 gr.Interface(detect_objects, inputs, outputs, title="YOLOv7 Object Detection").launch()