app.py

# Imports
import cv2
from detector import Detector
from pathlib import Path
import gradio as gr
import os
import numpy as np
import csv
import matplotlib.pyplot as plt
import pandas as pd

# Choose weights, names and config file
chosen_weights = "cfg/vessels_tyv4.weights"
chosen_config_file = "cfg/vessels_tyv4.cfg"
chosen_names = "cfg/vessel.names"

with open(chosen_names, "r") as f:
    classes = [line.strip() for line in f.readlines()]
labels = np.array(classes, dtype=str)

# Function for analyzing images
def analyze_image(selected_image, chosen_conf_thresh, chosen_nms_thresh):

    # Delete existing output files
    if os.path.exists("Classes.csv"):
        os.remove("Classes.csv")

    if selected_image is None:
        raise RuntimeError("No image found!")

    print("Starting image scan")
    
    # Initialize the detector
    detector = Detector(weights = str(chosen_weights), config_file = str(chosen_config_file), classes_file = chosen_names,  conf_thresh = chosen_conf_thresh, nms_thresh = chosen_nms_thresh)

    # Detect the image
    img_det, classes_id, scores, boxes = detector.detect(selected_image)

    class_names = []

    for _class in classes_id:
        class_names.append(labels[_class])

    tags = []

    for i in range(len(class_names)):
        tags.append([str(class_names[i]), str(boxes[i][0]), str(boxes[i][1]), str(boxes[i][2]), str(boxes[i][3]), str(scores[i])])
        
    print("Image scan finished succefully.")

    # Save tags in a csv file
    with open("Classes.csv", "w") as f:
        write = csv.writer(f)  
        write.writerow(["Class", "X", "Y", "Width", "Height", "Score"])      
        write.writerows(tags)
        f.close()
        
    df = pd.DataFrame(tags, columns = ["Class", "X", "Y", "Width", "Height", "Score"])

    return img_det, "Classes.csv", df.head(10)

# Function for analyzing video
def analyze_video(selected_video, chosen_conf_thresh, chosen_nms_thresh, start_sec, duration):

    # Delete existing output files
    if os.path.exists("demo_film.mp4"):
        os.remove("demo_film.mp4")
    if os.path.exists("output.mp4"):
        os.remove("output.mp4")
    if os.path.exists("Classes.csv"):
        os.remove("Classes.csv")

    if selected_video is None:
        raise RuntimeError("No video found!")

    print("Starting video scan")

    # Capture the video input
    video = cv2.VideoCapture(selected_video)
    ret, frame = video.read()
    if not ret: # Checking
        raise RuntimeError("Cannot read video stream!")
        
    # Calculate start and end frame 
    total_frames = round(video.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = video.get(cv2.CAP_PROP_FPS) 
    start_frame = round(start_sec * fps)
    end_frame = round(start_frame + (duration * fps))

    # Initialize the detector
    detector = Detector(weights = str(chosen_weights), config_file = str(chosen_config_file), classes_file = chosen_names,  conf_thresh = chosen_conf_thresh, nms_thresh = chosen_nms_thresh)

    frame_id = 0
    vid_out = []
    save_file_name = "demo_film.mp4"

    tags = []
    unique_objects = []
    
    if start_frame > total_frames:
        raise RuntimeError("Start second is out of bounds!")

    while True:
        # Read the image
        ret, frame = video.read()

        if not ret: # Error or end of stream heck
            break
        if frame is None: continue
        
        if  start_frame <= frame_id <= end_frame:
            class_names = []
            # Detect the image
            img_det, classes_id, scores, boxes = detector.detect(frame)

            for _class in classes_id:
                class_names.append(labels[_class])

                if unique_objects.count(labels[_class]) == 0:
                    unique_objects.append(labels[_class])

            for i in range(len(class_names)):
                tags.append([str(class_names[i]), str(boxes[i][0]), str(boxes[i][1]), str(boxes[i][2]), str(boxes[i][3]), str(scores[i]), str(frame_id)])

            # video writer
            if frame_id == start_frame:
                Height, Width = img_det.shape[:2]
                fps = video.get(cv2.CAP_PROP_FPS) if 15 < video.get(cv2.CAP_PROP_FPS) <= 30 else 15
                vid_out = cv2.VideoWriter(save_file_name, cv2.VideoWriter_fourcc(*"MP4V"), fps, (Width, Height)) 
            vid_out.write(img_det)

        if frame_id > end_frame:
            break

        frame_id += 1

    # Release videos
    video.release()
    vid_out.release()

    # Save tags in a csv file
    with open("Classes.csv", "w") as f:      
        write = csv.writer(f) 
        write.writerow(["Class", "X", "Y", "Width", "Height", "Score","Frame"])     
        write.writerows(tags)
        f.close()

    if end_frame > total_frames:
        end_frame = total_frames

    # Create graph
    plt.switch_backend("agg")
    fig = plt.figure()

    df = pd.DataFrame(tags, columns = ["Class", "X", "Y", "Width", "Height", "Score", "Frame"])

    # For every different object found, check how many times it appears in each frame
    for unique_object in unique_objects:
        object_array = df[df["Class"].str.fullmatch(unique_object)==True]
        obj_per_frame = []

        for i in range(start_frame, end_frame + 1):
            temp_array = []
            temp_array = object_array[object_array["Frame"].astype("str").str.fullmatch(str(i))==True]
            rows = temp_array.shape[0]
            obj_per_frame.append(rows)

        # Plot line graph for every individual object found
        plt.plot(list(range(start_frame, end_frame+1)), obj_per_frame, label = unique_object)

    plt.title("Objects per frame")
    plt.ylabel("Objects")
    plt.xlabel("Frame")
    plt.legend()

    print("Video scan finished succefully.")

    # Changes video fourcc to h264 so that it can be displayed in the browser
    os.system("ffmpeg -i demo_film.mp4 -vcodec libx264 -f mp4 output.mp4")

    return "output.mp4", fig, "Classes.csv", df.head(10)


# Dradio interfaces take mandatory parameters an input function, the input type(s) and output type(s)
# Demo is hosted on http://localhost:7860/

# Examples
image_examples = [
    ["examples/vessels.png", 0.20, 0.40],
    ["examples/boat.png", 0.25, 0.40],
]
video_examples =[["examples/vessels.mp4", 0.25, 0.40, 0, 10]]

# Image interface
image_interface = gr.Interface(fn = analyze_image,
                                   inputs = [gr.Image(label = "Image"),
                                   gr.Slider(0, 1, value = 0.25, label = "Confidence Threshold"),
                                   gr.Slider(0, 1, value = 0.40, label = "Non Maxima Supression threshold")],
                                   outputs = [gr.Image(label="Image"), gr.File(label="All classes"), gr.Dataframe(label="Ten first classes", headers = ["Class", "X", "Y", "Width", "Height", "Score"])],
                                   allow_flagging = False,
                                   cache_examples = False,
                                   examples = image_examples)

# Video interface
video_interface = gr.Interface(fn = analyze_video,
                                  inputs = [gr.Video(label = "Video"),
                                  gr.Slider(0, 1, value = 0.25, label = "Confidence Threshold"),
                                  gr.Slider(0, 1, value = 0.40, label = "Non Maxima Supression threshold"),
                                  gr.Slider(0, 60, value = 0, label = "Start Second", step = 1),
                                  gr.Slider(1, 10, value = 4, label = "Duration", step = 1)],
                                  outputs = [gr.Video(label="Video"), gr.Plot(label="Objects per frame"), gr.File(label="All classes"), gr.Dataframe(label=" Ten first classes", headers = ["Class", "X", "Y", "Width", "Height", "Score", "Frame"])],
                                  allow_flagging = False,
                                  cache_examples = False,
                                  examples = video_examples)

gr.TabbedInterface(
    [video_interface, image_interface],
    ["Scan Videos", "Scan Images"]
).launch()