ssim-tests

ssim-tests.py

import io
import os
import pprint

import boto3
import cv2
import numpy as np
from PIL import Image, ImageOps
from skimage import img_as_float
from skimage.measure import compare_ssim as ssim
s3 = boto3.client('s3')
rekog = boto3.client('rekognition')

#cap = cv2.VideoCapture(os.path.expanduser('~/Downloads/test.mp4'))
cap = cv2.VideoCapture(os.path.expanduser('~/Downloads/Standardized Testing.mp4'))

def detect_new_frames(cap, seconds, threshold):
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    retval, old_frame = cap.read()
    old_frame = img_as_float(cv2.cvtColor(old_frame, cv2.COLOR_BGR2GRAY))
    old_frame = img_as_float(old_frame)
    frames = 0
    if not retval:
        return
    while True:
        retval, frame = cap.read()
        frames += 1
        if not retval:
            break
        if frames % (seconds*fps) == 0:
            frame_gray = img_as_float(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY))
            ssim_v = ssim(
                old_frame,
                frame_gray,
                data_range=frame.max() - frame.min()
            )
            old_frame = frame_gray
            if ssim_v < threshold:
                yield cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

def get_labels(frames):
    for frame in frames:
        im = Image.fromarray(frame)
        buff = io.BytesIO()
        im.save(buff, "JPEG")
        buff.seek(0)
        labels = rekog.detect_labels(Image={'Bytes': buff.read()})['Labels']
        yield labels

pprint.pprint(list(get_labels(detect_new_frames(cap, 6, 0.8))))