賴竹維
賴竹維

Reputation: 21

How to print the class name from yolo object detection?

I am new to using the YOLOv4 for a real-time object detection application for a project.

I followed the steps from the darknet website. And for the real-time detection using a webcam, I ran this code below

import cv2
import numpy as np
import time

# Load the YOLO model
net = cv2.dnn.readNet("./solarpanel/cfg/weights/yolov4-tiny-custom_last.weights",
                      "./solarpanel/cfg/yolov4-tiny-custom.cfg")
classes = []
with open("./solarpanel/cfg/obj.names", "r") as f:
    classes = [line.strip() for line in f.readlines()]
layer_names = net.getLayerNames()
output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]
colors = np.random.uniform(0, 255, size=(len(classes), 3))

cap=None
font=None
starting_time = None
frame_id = 0
def loadWebCam():
    global cap, font, starting_time, frame_id
    # Load webcam
    cap = cv2.VideoCapture(0)
    font = cv2.FONT_HERSHEY_SIMPLEX
    starting_time = time.time()
    frame_id = 0
    

def unLoadWebCam():
    global cap
    cap.release()
    cv2.destroyAllWindows()


def checkcellfailure():
    global cap, font, starting_time, frame_id
    loadWebCam()
    class_ids = []
    confidences = []
    boxes = []
    while True:
        # Read webcam
        _, frame = cap.read()
        frame_id += 1
        height, width, channels = frame.shape

        # Detecting objects
        blob = cv2.dnn.blobFromImage(frame,
                                    0.00392, (416, 416), (0, 0, 0),
                                    True,
                                    crop=False)
        net.setInput(blob)
        outs = net.forward(output_layers)

        # Visualising data
        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.1:
                    # Object detected
                    center_x = int(detection[0] * width)
                    center_y = int(detection[1] * height)
                    w = int(detection[2] * width)
                    h = int(detection[3] * height)

                    # Rectangle coordinates
                    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)

        indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.8, 0.3)

        for i in range(len(boxes)):
            if i in indexes:
                x, y, w, h = boxes[i]
                label = str(classes[class_ids[i]])
                confidence = confidences[i]
                color = colors[class_ids[i]]
                cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
                cv2.putText(frame, label + " " + str(round(confidence, 2)),
                            (x, y + 30), font, 3, color, 3)

        elapsed_time = time.time() - starting_time
        fps = frame_id / elapsed_time
        cv2.putText(frame, "FPS: " + str(round(fps, 2)), (40, 670), font, .7,
                    (0, 255, 255), 1)
        cv2.putText(frame, "press [esc] to exit", (40, 690), font, .45,
                    (0, 255, 255), 1)
        cv2.imshow("Image", frame)
        key = cv2.waitKey(1)
       # print(confidences)
        found=False
        for conf in confidences:
            if conf > 0.9:
                found=True
        if key == 27 or found:
            print("cellfailure detected")
            #print(" Global Location: %s" % vehicle.location.global_frame)
            #print(" Global Location (relative altitude): %s" % vehicle.location.global_relative_frame)
            #print(" Local Location: %s" % vehicle.location.local_frame)
            #print("[button pressed] ///// [esc].")
            #print("[feedback] ///// Videocapturing succesfully stopped")
            break
    #unLoadWebCam()
    return confidences

This opens up a new camera window that detects the classes and closes the webcam.(wherein classes I have 2 classes) but I need only specified classes as output.

Could anyone suggest me the right way of code to print the specified class as output and to print the location of the drone?

Upvotes: 2

Views: 2339

Answers (0)

Related Questions