Update app.py

app.py

@@ -1,141 +1,83 @@
 import gradio as gr
-import spaces
-from huggingface_hub import hf_hub_download
+import cv2
+import numpy as np
 
-def download_models(model_id):
-    hf_hub_download("qbcode/chicken_detection", filename=f"{model_id}", local_dir=f"./")
-    return f"./{model_id}"
-
-@spaces.GPU
-def yolov8_inference(img_path, model_id, image_size, conf_threshold, iou_threshold):
-    """
-    Load a YOLOv9 model, configure it, perform inference on an image, and optionally adjust 
-    the input size and apply test time augmentation.
-    
-    :param model_path: Path to the YOLOv9 model file.
-    :param conf_threshold: Confidence threshold for NMS.
-    :param iou_threshold: IoU threshold for NMS.
-    :param img_path: Path to the image file.
-    :param size: Optional, input size for inference.
-    :return: Annotated image with the number of objects detected.
-    """
-    # Import YOLOv9
-    import yolov8
-    
-    
-    # Load the model
-    model_path = download_models(model_id)
-    model = yolov8.load(model_path, device="cuda:0")
+# 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"
     
-    # Set model parameters
-    model.conf = conf_threshold
-    model.iou = iou_threshold
-    
-    # Perform inference
-    results = model(img_path, size=image_size)
+    # 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)
 
-    # Optionally, show detection bounding boxes on image
-    annotated_image = results.render()
-    
-    # Count the number of detected objects
-    num_objects_detected = len(results.pred[0])
+    # 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)
     
-    return annotated_image, f"Number of objects detected: {num_objects_detected}"
+    # 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)
 
-def app():
-    with gr.Blocks():
-        with gr.Row():
-            with gr.Column():
-                img_path = gr.Image(type="filepath", label="Image")
-                model_path = gr.Dropdown(
-                    label="Model",
-                    choices=[
-                        "best.pt",
-                        "gelan-e.pt",
-                        "yolov9-c.pt",
-                        "yolov9-e.pt",
-                    ],
-                    value="gelan-e.pt",
-                )
-                image_size = gr.Slider(
-                    label="Image Size",
-                    minimum=320,
-                    maximum=1280,
-                    step=32,
-                    value=640,
-                )
-                conf_threshold = gr.Slider(
-                    label="Confidence Threshold",
-                    minimum=0.1,
-                    maximum=1.0,
-                    step=0.1,
-                    value=0.4,
-                )
-                iou_threshold = gr.Slider(
-                    label="IoU Threshold",
-                    minimum=0.1,
-                    maximum=1.0,
-                    step=0.1,
-                    value=0.5,
-                )
-                yolov8_infer = gr.Button(value="Inference")
+    return image
 
-            with gr.Column():
-                output_image = gr.Image(type="numpy",label="Output")
-                num_objects_detected = gr.Textbox(label="Number of Objects Detected")
+# Create Gradio interface
+models = ["YOLOv7", "YOLOv8", "YOLO-NAS"]
 
-        yolov8_infer.click(
-            fn=yolov8_inference,
-            inputs=[
-                img_path,
-                model_path,
-                image_size,
-                conf_threshold,
-                iou_threshold,
-            ],
-            outputs=[output_image, num_objects_detected],
-        )
-        
-        gr.Examples(
-            examples=[
-                [
-                    "eg/test1",
-                    "best.pt",
-                    640,
-                    0.4,
-                    0.5,
-                ],
-            ],
-            fn=yolov8_inference,
-            inputs=[
-                img_path,
-                model_path,
-                image_size,
-                conf_threshold,
-                iou_threshold,
-            ],
-            outputs=[output_image, num_objects_detected],
-            cache_examples=True,
-        )
+pages = []
+for model in models:
+    inputs = [
+        gr.inputs.Image(label="Upload Image or Video"),
+    ]
 
+    outputs = gr.outputs.Image(label=f"{model} Output")
 
-gradio_app = gr.Blocks()
-with gradio_app:
-    gr.HTML(
-        """
-    <h1 style='text-align: center'>
-    YOLOv9: Learning What You Want to Learn Using Programmable Gradient Information
-    </h1>
-    """)
-    gr.HTML(
-        """
-        <h3 style='text-align: center'>
-        Follow me for more!
-        <a href='https://twitter.com/kadirnar_ai' target='_blank'>Twitter</a> | <a href='https://github.com/kadirnar' target='_blank'>Github</a> | <a href='https://www.linkedin.com/in/kadir-nar/' target='_blank'>Linkedin</a>  | <a href='https://www.huggingface.co/kadirnar/' target='_blank'>HuggingFace</a>
-        </h3>
-        """)
-    with gr.Row():
-        with gr.Column():
-            app()
+    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))
 
-gradio_app.launch(debug=True)
+gr.Interface(pages, gr.Interface.MULTIPAGE).launch()