analysis.py

import cv2
import os
from matplotlib import pyplot as plt
import numpy as np

def maxrect(fig):
    lower_mask = np.array([45, 25, 153])
    upper_mask = np.array([60, 62, 205])
    hsv = cv2.cvtColor(fig, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(hsv, lower_mask, upper_mask)
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
    areas = []
    for c in range(len(contours)):
        areas.append(cv2.contourArea(contours[c]))
    max_id = areas.index(max(areas))
    x, y, w, h = cv2.boundingRect(contours[max_id])
    fig = cv2.drawContours(fig, contours, -1, (255, 255, 255), 4)
    fig = cv2.rectangle(fig, (x, y), (x + w, y + h), (0, 0, 255), 4)
    return fig

path = 'dataset'
files = os.listdir(path)

figure, axes = plt.subplots(6, 5)
figure.set_figheight(8)
figure.set_figwidth(11)

for i, file in enumerate(files):
    seg = cv2.VideoCapture(path + '/' + file)
    frame_set = []
    print(f'Processing video {i+1}')
    while True:
        ret, frame = seg.read()
        if not ret:
            break
        frame = cv2.flip(frame, 1)
        frame_set.append(frame)

    fig = frame_set[len(frame_set)//2]
    fig = maxrect(fig)
    
    fig = cv2.cvtColor(fig, cv2.COLOR_BGR2RGB)
    axes[i//5, i%5].imshow(fig)

for ax in axes.flat:
    ax.set_xticks([])
    ax.set_yticks([])
    for spine in ax.spines.values():
        spine.set_visible(False)
for i, ax in enumerate(axes[:, 0]):
    ax.set_ylabel(f'm={i+1}')
for i, ax in enumerate(axes[-1, :]):
    ax.set_xlabel(f'd={5-i}')

plt.tight_layout()
plt.savefig('multi-1.pdf', dpi=300)
plt.show()