Update app.py

app.py

@@ -1,56 +1,34 @@
-import gradio as gr
-import cv2
+import torch
+from torchvision import transforms
 import numpy as np
 
 # Load YOLOv7 model
 weights_path = "cattle.pt"
 config_path = "yolov7.yaml"
-net = cv2.dnn.readNet(weights_path, config_path)
-layer_names = net.getLayerNames()
-output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
+net = None  # Initialize your YOLOv7 model here
 
 # Define function to detect objects using YOLOv7
 def detect_objects(image):
-    # Detect objects in the image
-    blob = cv2.dnn.blobFromImage(image, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
-    net.setInput(blob)
-    outs = net.forward(output_layers)
-
-    # Process detection results
-    class_ids = []
-    confidences = []
-    boxes = []
-    for out in outs:
-        for detection in out:
-            scores = detection[5:]
-            class_id = np.argmax(scores)
-            confidence = scores[class_id]
-            if confidence > 0.5:
-                # Object detected
-                center_x = int(detection[0] * image.shape[1])
-                center_y = int(detection[1] * image.shape[0])
-                w = int(detection[2] * image.shape[1])
-                h = int(detection[3] * image.shape[0])
-                x = int(center_x - w / 2)
-                y = int(center_y - h / 2)
-                boxes.append([x, y, w, h])
-                confidences.append(float(confidence))
-                class_ids.append(class_id)
-
-    # Non-max suppression to remove overlapping boxes
-    indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
-
-    # Count detected objects
-    object_count = len(indexes)
-
-    # Draw bounding boxes on the image
-    for i in range(len(boxes)):
-        if i in indexes:
-            x, y, w, h = boxes[i]
-            label = str(class_ids[i])
-            color = (255, 0, 0)  # BGR color format
-            cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
-            cv2.putText(image, label, (x, y + 30), cv2.FONT_HERSHEY_PLAIN, 3, color, 3)
+    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