Files
2024-08-02 19:39:19 -06:00

82 lines
2.0 KiB
Python

import os
import sys
import cv2
import numpy as np
MINIMUM_SIMILARITY_THRESHOLD = 0.4
def process_video(video_path, min_similarity, output_folder):
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
print("Error opening video file")
return
prev_frame = None
kp1, ds1 = None, None
i = 0
while cap.isOpened():
ret, frame = cap.read()
if not ret: break
i += 1
print(f"Frame {i}: ", end='')
if prev_frame is None:
kp1, des1 = orb.detectAndCompute(frame, None)
cv2.imwrite(os.path.join(output_folder, os.path.basename(video_path)+"-"+str(i)+".png"), frame)
print("Initial frame")
prev_frame = frame
continue
cv2.imshow('Frame', frame)
similarity = orb_similarity(kp1, des1, frame)
print(f"Similarity: {similarity}, ", end='')
if similarity < min_similarity:
kp1, des1 = orb.detectAndCompute(prev_frame, None)
if(not os.path.isdir(output_folder)):
os.mkdir(output_folder)
cv2.imwrite(os.path.join(output_folder, os.path.basename(video_path)+"-"+str(i)+".png"), prev_frame)
print("Lost continuity, new match frame")
else:
print("Skipped!")
prev_frame = frame
if cv2.waitKey(1) == ord('q'): break
# Release the video capture object and close windows
cap.release()
cv2.destroyAllWindows()
orb = cv2.ORB_create()
bf = cv2.BFMatcher(cv2.NORM_L1, crossCheck=True)
def orb_similarity(kp1, des1, img2):
kp2, des2 = orb.detectAndCompute(img2, None)
matches = bf.match(des1, des2)
# try:
# matches = bf.match(des1, des2)
# except:
# return 0.0
similarity = len(matches) / min(len(kp1), len(kp2))
return similarity
if len(sys.argv) != 4:
print("Usage: python main.py <video file> <min similarity threshold> <output folder>")
else:
process_video(sys.argv[1], float(sys.argv[2]), sys.argv[3])