CaesarObjectDetection/CaesarYolo.py

import cv2
import numpy as np
import os
import time
import ffmpeg


class CaesarYolo:
    def __init__(self) -> None:
        self.CONFIDENCE = 0.5
        self.SCORE_THRESHOLD = 0.5
        self.IOU_THRESHOLD = 0.5
        self.current_dir = os.path.realpath(__file__).replace(f"/CaesarYolo.py","")
        config_path = f"{self.current_dir}/cfg/yolov3.cfg"
        weights_path = f"{self.current_dir}/weights/yolov3.weights"
        self.font_scale = 1
        self.thickness = 1
        self.LABELS = open(f"{self.current_dir}/data/coco.names").read().strip().split("\n")
        self.COLORS = np.random.randint(0, 255, size=(len(self.LABELS), 3), dtype="uint8")

        self.net = cv2.dnn.readNetFromDarknet(config_path, weights_path)

        self.ln = self.net.getLayerNames()
        try:
            self.ln = [self.ln[i[0] - 1] for i in self.net.getUnconnectedOutLayers()]
        except IndexError:
            # in case getUnconnectedOutLayers() returns 1D array when CUDA isn't available
            self.ln = [self.ln[i - 1] for i in self.net.getUnconnectedOutLayers()]
    @staticmethod
    def compress_video(video_full_path, output_file_name, target_size):
        # Reference: https://en.wikipedia.org/wiki/Bit_rate#Encoding_bit_rate
        min_audio_bitrate = 32000
        max_audio_bitrate = 256000

        probe = ffmpeg.probe(video_full_path)
        # Video duration, in s.
        duration = float(probe['format']['duration'])
        # Audio bitrate, in bps.
        audio_bitrate = float(next((s for s in probe['streams'] if s['codec_type'] == 'audio'), None)['bit_rate'])
        # Target total bitrate, in bps.
        target_total_bitrate = (target_size * 1024 * 8) / (1.073741824 * duration)

        # Target audio bitrate, in bps
        if 10 * audio_bitrate > target_total_bitrate:
            audio_bitrate = target_total_bitrate / 10
            if audio_bitrate < min_audio_bitrate < target_total_bitrate:
                audio_bitrate = min_audio_bitrate
            elif audio_bitrate > max_audio_bitrate:
                audio_bitrate = max_audio_bitrate
        # Target video bitrate, in bps.
        video_bitrate = target_total_bitrate - audio_bitrate

        i = ffmpeg.input(video_full_path)
        ffmpeg.output(i, os.devnull,
                    **{'c:v': 'libx264', 'b:v': video_bitrate, 'pass': 1, 'f': 'mp4'}
                    ).overwrite_output().run()
        ffmpeg.output(i, output_file_name,
                    **{'c:v': 'libx264', 'b:v': video_bitrate, 'pass': 2, 'c:a': 'aac', 'b:a': audio_bitrate}
                    ).overwrite_output().run()


    def video_load(self,videofile):
        self.video_file = f"{ self.current_dir}/{videofile}"
        if self.video_file:
            self.cap = cv2.VideoCapture(self.video_file)
            _, image = self.cap.read()
            h, w = image.shape[:2]
            fourcc = cv2.VideoWriter_fourcc(*"XVID")
            frames = self.cap.get(cv2.CAP_PROP_FRAME_COUNT)
            fps = self.cap.get(cv2.CAP_PROP_FPS)
            
            # calculate duration of the video
            self.duration_seconds = round(frames / fps)
            self.out = cv2.VideoWriter(f"{self.current_dir}/output.avi", fourcc, 20.0, (w, h))
            self.overall_time_taken = []
    def caesar_object_detect(self,image,verbose=False):
        if self.video_file and image is "video":
            _,image = self.cap.read()
        try:
            h, w = image.shape[:2]
        except AttributeError as aex:
            return None,None,None
        blob = cv2.dnn.blobFromImage(image, 1/255.0, (416, 416), swapRB=True, crop=False)
        self.net.setInput(blob)
        start = time.perf_counter()
        layer_outputs = self.net.forward(self.ln)
        time_took = time.perf_counter() - start
        if verbose == True:
            print("Time took:", time_took)
        if self.video_file:
            self.overall_time_taken.append(time_took)
            time_elapsed = round(sum(self.overall_time_taken),3)
            approx_finish = self.duration_seconds *4.6 # seconds
        boxes, confidences, class_ids = [], [], []

        # loop over each of the layer outputs
        for output in layer_outputs:
            # loop over each of the object detections
            for detection in output:
                # extract the class id (label) and confidence (as a probability) of
                # the current object detection
                scores = detection[5:]
                class_id = np.argmax(scores)
                confidence = scores[class_id]
                # discard weak predictions by ensuring the detected
                # probability is greater than the minimum probability
                if confidence > self.CONFIDENCE:
                    # scale the bounding box coordinates back relative to the
                    # size of the image, keeping in mind that YOLO actually
                    # returns the center (x, y)-coordinates of the bounding
                    # box followed by the boxes' width and height
                    box = detection[:4] * np.array([w, h, w, h])
                    (centerX, centerY, width, height) = box.astype("int")

                    # use the center (x, y)-coordinates to derive the top and
                    # and left corner of the bounding box
                    x = int(centerX - (width / 2))
                    y = int(centerY - (height / 2))

                    # update our list of bounding box coordinates, confidences,
                    # and class IDs
                    boxes.append([x, y, int(width), int(height)])
                    confidences.append(float(confidence))
                    class_ids.append(class_id)

        # perform the non maximum suppression given the scores defined before
        idxs = cv2.dnn.NMSBoxes(boxes, confidences, self.SCORE_THRESHOLD, self.IOU_THRESHOLD)

        self.font_scale = 1
        self.thickness = 1

        # ensure at least one detection exists
        if len(idxs) > 0:
            # loop over the indexes we are keeping
            for i in idxs.flatten():
                # extract the bounding box coordinates
                x, y = boxes[i][0], boxes[i][1]
                w, h = boxes[i][2], boxes[i][3]
                # draw a bounding box rectangle and label on the image

                color = [int(c) for c in self.COLORS[class_ids[i]]]
                cv2.rectangle(image, (x, y), (x + w, y + h), color=color, thickness=self.thickness)
                text = f"{self.LABELS[class_ids[i]]}: {confidences[i]:.2f}"
                # calculate text width & height to draw the transparent boxes as background of the text
                (text_width, text_height) = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, fontScale=self.font_scale, thickness=self.thickness)[0]
                text_offset_x = x
                text_offset_y = y - 5
                box_coords = ((text_offset_x, text_offset_y), (text_offset_x + text_width + 2, text_offset_y - text_height))
                overlay = image.copy()
                cv2.rectangle(overlay, box_coords[0], box_coords[1], color=color, thickness=cv2.FILLED)
                # add opacity (transparency to the box)
                image = cv2.addWeighted(overlay, 0.6, image, 0.4, 0)
                # now put the text (label: confidence %)
                cv2.putText(image, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX,
                    fontScale=self.font_scale, color=(0, 0, 0), thickness=self.thickness)
        if self.video_file:
            self.out.write(image)
            return image,time_elapsed,approx_finish
        elif not self.video_file:
            return image,0,0

if __name__ == "__main__":
    def test():
        caesaryolo = CaesarYolo()
        caesaryolo.video_load("car-detection.mp4")
        while True:
            image,time_elapsed,end_time = caesaryolo.caesar_object_detect("video")
            if image is not None:
                print(round(time_elapsed,3),"out of",end_time)
                cv2.imshow("image", image)
            
                if ord("q") == cv2.waitKey(1):
                    break
            else:
                break


        caesaryolo.cap.release()
        cv2.destroyAllWindows()
    def convert_avi_to_mp4(avi_file_path, output_name):
        os.system(f"ffmpeg -y -i {avi_file_path} {output_name}") 
        return True

    CURRENT_DIR = os.path.realpath(__file__).replace(f"/CaesarYolo.py","")
    #convert_avi_to_mp4(,)
    import subprocess

    #process = subprocess.Popen(ffmpeg_command, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, bufsize = -1)