Update app.py

app.py

@@ -2,30 +2,15 @@ import gradio as gr
 import cv2
 import numpy as np
 
-# Define function to detect objects using YOLO
-def detect_objects(image, model_type):
-    # Load the YOLO model based on the model type
-    if model_type == 'YOLOv7':
-        weights_path = "path_to_yolov7_weights"
-        config_path = "path_to_yolov7_config"
-        net = cv2.dnn.readNet(weights_path, config_path)
-        layer_names = net.getLayerNames()
-        output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
-    elif model_type == 'YOLOv8':
-        weights_path = "path_to_yolov8_weights"
-        config_path = "path_to_yolov8_config"
-        net = cv2.dnn.readNet(weights_path, config_path)
-        layer_names = net.getLayerNames()
-        output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
-    elif model_type == 'YOLO-NAS':
-        weights_path = "path_to_yolo_nas_weights"
-        config_path = "path_to_yolo_nas_config"
-        net = cv2.dnn.readNet(weights_path, config_path)
-        layer_names = net.getLayerNames()
-        output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
-    else:
-        return "Invalid Model Type"
-    
+# 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()]
+
+# 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)
@@ -54,7 +39,10 @@ def detect_objects(image, model_type):
 
     # 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:
@@ -64,20 +52,10 @@ def detect_objects(image, model_type):
             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)
 
-    return image
+    return image, object_count
 
 # Create Gradio interface
-models = ["YOLOv7", "YOLOv8", "YOLO-NAS"]
-
-pages = []
-for model in models:
-    inputs = [
-        gr.inputs.Image(label="Upload Image or Video"),
-    ]
-
-    outputs = gr.outputs.Image(label=f"{model} Output")
-
-    page = gr.Interface(detect_objects, inputs, outputs, title=f"{model} Object Detection", description=f"Identify objects in images or videos using {model}").launch()
-    pages.append((f"{model} Model", page))
+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(pages, gr.Interface.MULTIPAGE).launch()
+gr.Interface(detect_objects, inputs, outputs, title="YOLOv7 Object Detection").launch()