I have folder with short videos and folder with images. Most of the images are sceenshots from one of the videos, but they may not be exactly the same (different size, noise, loss of details because of compression, etc). My goal is to match every image with video it was taken from. So far, I use OpenCV library to load one video and calculate SSIM score between each video frame and each image. I store the highest SSIM score of every image. Then I'd take the image with highest SSIM score, associate it with the video and run the function again for second video.
Here is my code:
import cv2
import numpy as np
from skimage.measure import compare_ssim
import sqlite3
#screenshots - list that contains dict(id=screenshot id, image=jpeg image data)
#video_file - str - path to video file
def generate_matches(screenshots, video_file):
for screenshot in screenshots:
screenshot["cv_img"] = cv2.imdecode(np.fromstring(screenshot["image"], np.uint8), 0)
screenshot["best_match"] = dict(score=0, frame=0)
screenshot.pop('image', None) #remove jpg data from RAM
vidcap = cv2.VideoCapture(video_file)
success,image = vidcap.read()
count = 1
while success:
image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
for screenshot in screenshots:
c_image = cv2.resize(image, screenshot["cv_img"].shape[1::-1])
score = compare_ssim(screenshot["cv_img"], c_image, full=False)
if score > screenshot["best_match"]["score"]:
screenshot["best_match"] = dict(score=score,frame=count)
count += 1
success,image = vidcap.read()
if count % 500 == 0:
print("Frame {}".format(count))
print("Last Frame {}".format(count))
for screenshot in screenshots:
c.execute("INSERT INTO matches(screenshot_id, file, match, frame) VALUE (?,?,?,?)",
(screenshot["id"], video_file, screenshot["best_match"]["score"], screenshot["best_match"]["frame"]))
generate_matches(list_of_screenshots, "video1.mp4")
generate_matches(list_of_screenshots, "video2.mp4")
...
This algorithm seems to be good-enough to associate videos with images, but it is quite slow, even if I'd use more threads. Is there any way to make it faster? Maybe different algorithm or some pre-processing of videos and images? I'll be glad for any ideas!