Skip to content

Instantly share code, notes, and snippets.

@adhorn
Forked from ranman/ssim-tests.py
Created August 4, 2017 15:24
Show Gist options
  • Save adhorn/09490dff06e33fb8a9527f8f1e4a7c2d to your computer and use it in GitHub Desktop.
Save adhorn/09490dff06e33fb8a9527f8f1e4a7c2d to your computer and use it in GitHub Desktop.
ssim-tests
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))))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment