import cv2
import numpy as np
from pathlib import Path
import time
from tqdm import tqdm
import argparse
import pandas as pd
P_HASH = cv2.img_hash.PHash_create()
POP_COUNT = np.array([bin(i).count("1") for i in range(256)], dtype=np.uint8)
def p_hash_bytes(Gray_img):
h = P_HASH.compute(Gray_img)
return h.reshape(-1).copy()
def grid_phash(gray_img, grid=(3,3)):
gh, gw = grid
H, W = gray_img.shape
th, tw = H // gh, W // gw
parts = []
for r in range(gh):
for c in range (gh):
tile = gray_img[r*th:(r+1)*th, c*tw:(c+1)*tw]
parts.append(p_hash_bytes(tile))
return np.concatenate(parts, axis=0)
def hamming(a, b):
return int(POP_COUNT[np.bitwise_xor(a, b)].sum())
def mse(img1, img2):
h, w = img1.shape
diff = cv2.subtract(img1, img2)
err = np.sum(diff**2)
mse = err/(float(h*w))
return mse
def centerCrop(img:np, size:tuple = None, pct_size:float = 0.3):
if size is None:
size = int(img.shape[0] * (1 - pct_size)), int(img.shape[1] * (1 - pct_size))
x, y = (img.shape[1] - size[1]) // 2, (img.shape[0] - size[0]) // 2
return img[y:y+size[0], x:x+size[1]]
def removeBlackBorder(image):
copyImg = cv2.cvtColor(image.copy(), cv2.COLOR_BGR2HSV)
h = copyImg[:,:,0]
mask = np.ones(h.shape, dtype=np.uint8) * 255
th = (25, 175)
mask[(h > th[0]) & (h < th[1])] = 0
copyImg = cv2.cvtColor(copyImg, cv2.COLOR_HSV2BGR)
resROI = cv2.bitwise_and(copyImg, copyImg, mask=mask)
image_gray = cv2.cvtColor(resROI, cv2.COLOR_BGR2GRAY)
_, thresh = cv2.threshold(image_gray, 0, 255, cv2.THRESH_BINARY)
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))
morph = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
contours = cv2.findContours(morph, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
contours = contours[0] if len(contours) == 2 else contours[1]
bigCont = max(contours, key=cv2.contourArea)
x, y, w, h = cv2.boundingRect(bigCont)
crop = image[y : y + h, x : x + w]
return crop, (x, y, w, h)
def normalization(img, mean = 0.458971, std = 0.225609):
Norm_img = cv2.cvtColor(img.copy(), cv2.COLOR_BGR2GRAY).astype(np.float32) / 255.0
Norm_img = (Norm_img - mean) / std
Norm_img = cv2.GaussianBlur(Norm_img, (5, 5), 1)
return Norm_img
def seconds_to_hms(seconds):
hours = seconds // 3600
minutes = (seconds % 3600) // 60
secs = seconds % 60
return f'{int(hours)}:{int(minutes):02}:{int(secs):02}'
def fft_cross_correlation(img1, img2):
h, w = img1.shape
H, W = cv2.getOptimalDFTSize(2*h), cv2.getOptimalDFTSize(2*w)
pad1 = np.zeros((H, W), dtype=np.float32)
pad2 = np.zeros((H, W), dtype=np.float32)
pad1[:h, :w] = img1
pad2[:h, :w] = img2
dtf1 = cv2.dft(pad1, flags=cv2.DFT_COMPLEX_OUTPUT)
dtf2 = cv2.dft(pad2, flags=cv2.DFT_COMPLEX_OUTPUT)
cross_power = cv2.mulSpectrums(dtf1, dtf2, 0, conjB=True)
corr = cv2.idft(cross_power, flags=cv2.DFT_SCALE | cv2.DFT_REAL_OUTPUT)
corr = np.fft.fftshift(corr)
max_val = corr.max()
n_img1 = np.linalg.norm(img1)
n_img2 = np.linalg.norm(img2)
sim = max_val / (n_img1 * n_img2)
return corr, sim
def find_frame_in_video(Target, Input, Patience:int = 60, Remove_Blackbar:bool = False, Forze_resize:bool = False, Sim_Check:bool = True):
Target = Path(Target)
Input = Path(Input)
grid_size = (3,3)
start_time = time.time()
img = cv2.imread(str(Target), cv2.IMREAD_COLOR)
if Remove_Blackbar:
img, _ = removeBlackBorder(img)
#img = normalization(img)
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
query_img_hash = grid_phash(img, grid_size)
cap = cv2.VideoCapture(str(Input))
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
video_fps = cap.get(cv2.CAP_PROP_FPS)
print(f"{Input.name} Video FPS: {video_fps:.2f}, Total frames: {total_frames}, Video length: {seconds_to_hms(total_frames / video_fps)}")
pbar = tqdm(
total=total_frames,
bar_format="{desc}", # only render the description
desc="aFPS: 0.00" # initial text
)
best_match_segs = []
prev = time.perf_counter()
best_match = 10000
frameCouter = 0
best_match_frame = 0
ROI_rect = None
resizeFrame = False
Max_Sim = 0.90
early_Stop_count = 0
sim = 0
sim_str = ""
prev_sim = 0
segCout = 0
while cap.isOpened():
if early_Stop_count == Patience:
break
#cap.set(cv2.CAP_PROP_POS_FRAMES, frameCouter)
cap.set(cv2.CAP_PROP_POS_MSEC, segCout * 1000)
ret, frame = cap.read()
if not ret:
break
try:
if (frameCouter == 0):
if Remove_Blackbar:
frame, ROI_rect = removeBlackBorder(frame)
if Remove_Blackbar and frameCouter > 0:
x, y, w, h = ROI_rect
frame = frame[y : y + h, x : x + w]
#frame = normalization(frame)
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
frame_hash = grid_phash(frame, grid_size)
sim = hamming(query_img_hash, frame_hash)
if sim < best_match:
best_match = sim
best_match_frame = frameCouter
#best_match_segs.append(best_match_frame / video_fps)
best_match_segs.append(segCout)
early_Stop_count += 1 if (Patience is not None) and (sim >= Max_Sim) else 0
except Exception as e:
continue
finally:
now = time.perf_counter()
fps = 1.0 / (now - prev)
prev = now
#pbar.set_description(f"Progress: {(frameCouter * 100 / total_frames):.2f}% {frameCouter} / {total_frames}, Time: {seconds_to_hms(frameCouter / video_fps)} FPS: {fps:.2f} best match {best_match:.6f}/{sim:.6f} Aprox timestamp {seconds_to_hms(best_match_frame / video_fps)}")
pbar.set_description(f"Progress: {(frameCouter * 100 / total_frames):.2f}% {frameCouter} / {total_frames}, Time: {seconds_to_hms(segCout)} FPS: {fps:.2f} best match {best_match:.6f}/{sim:.6f} Aprox timestamp {seconds_to_hms(best_match_frame / video_fps)}")
pbar.update(1)
frameCouter += round(video_fps)
segCout += 1
pbar.close()
cap.release()
end_time = time.time()
print(f"{seconds_to_hms(end_time - start_time)} seconds")
return best_match_segs[-1], 0, best_match
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("-f", "--File", type=str, required=True)
parser.add_argument("-o", "--Output", type=str, default="output.csv")
parser.add_argument("--Sim_Check", type=bool, default=False)
parser.add_argument("-p", "--Patience", type=int, default=None, help="Seconds analyzed after the threshold is reached")
parser.add_argument("-br","--Remove_Blackbar", type=bool, default=False)
parser.add_argument("-r","--Forze_resize", type=bool, default=False)
args = parser.parse_args()
outputFile = Path(args.Output)
file_csv = pd.read_csv(Path(args.File))
if outputFile.exists():
output_file = pd.read_csv(outputFile)
else:
output_file = pd.DataFrame({
"case": [0],
"video_path": [""],
"image_path": [""],
"best_frame_seg": [0],
"error": [0],
"sim": [0]
})
for i, row in file_csv.iterrows():
if pd.isnull(row["image_path"]):
print(f"case {row['case']} skipped")
continue
videos = Path(row["video_path"])
if videos.is_file():
list_videos = [videos]
else:
list_videos = sorted(list(videos.glob("**/*.[mM][pP]4")))
print(f"case {row['case']}")
for videos in list_videos:
try:
res = find_frame_in_video(row["image_path"], videos, args.Patience, args.Remove_Blackbar, args.Forze_resize, args.Sim_Check)
new_row = pd.DataFrame({
"case": [row["case"]],
"video_path": [videos],
"image_path": [row["image_path"]],
"best_frame_seg": [res[0]],
"error": [res[1]],
"sim": [res[2]]
})
output_file = pd.concat([output_file, new_row], ignore_index = True)
except Exception as e:
print (e)
continue
output_file.to_csv(Path(args.Output), index=False)
