151 lines
3.2 KiB
Python
151 lines
3.2 KiB
Python
import subprocess
|
|
import numpy as np
|
|
import cv2
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from scipy import signal
|
|
|
|
TARGET_FPS = 30
|
|
|
|
def get_device():
|
|
|
|
if torch.cuda.is_available():
|
|
device = torch.device("cuda")
|
|
|
|
print("GPU detected")
|
|
print("Device:", torch.cuda.get_device_name(0))
|
|
print("ROCm build:", torch.version.hip is not None)
|
|
|
|
return device
|
|
|
|
print("GPU not available, falling back to CPU")
|
|
return torch.device("cpu")
|
|
|
|
|
|
device = get_device()
|
|
|
|
|
|
def ffmpeg_frame_stream(path):
|
|
|
|
cmd = [
|
|
"ffmpeg",
|
|
"-i", path,
|
|
"-vf", "fps=30,scale=64:36,format=bgr24",
|
|
"-f", "rawvideo",
|
|
"-pix_fmt", "bgr24",
|
|
"-"
|
|
]
|
|
|
|
proc = subprocess.Popen(
|
|
cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
bufsize=10**8
|
|
)
|
|
|
|
frame_size = 64 * 36 * 3
|
|
|
|
while True:
|
|
raw = proc.stdout.read(frame_size)
|
|
if len(raw) < frame_size:
|
|
break
|
|
|
|
frame = np.frombuffer(raw, np.uint8).reshape((36, 64, 3))
|
|
yield frame
|
|
|
|
proc.stdout.close()
|
|
proc.wait()
|
|
|
|
|
|
def compute_signal(video_path):
|
|
|
|
prev_blocks = None
|
|
signal_out = []
|
|
|
|
for frame in ffmpeg_frame_stream(video_path):
|
|
|
|
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
|
|
|
t = torch.from_numpy(gray).float().to(device) / 255.0
|
|
|
|
# reshape for conv
|
|
t = t.unsqueeze(0).unsqueeze(0)
|
|
|
|
# block averages using avg pool
|
|
blocks = F.avg_pool2d(t, kernel_size=8, stride=8)
|
|
|
|
blocks = blocks.squeeze()
|
|
|
|
if prev_blocks is None:
|
|
prev_blocks = blocks
|
|
signal_out.append(0)
|
|
continue
|
|
|
|
# block difference energy
|
|
diff = torch.abs(blocks - prev_blocks).sum()
|
|
|
|
# edge energy
|
|
sobel_x = torch.tensor(
|
|
[[-1,0,1],[-2,0,2],[-1,0,1]],
|
|
dtype=torch.float32, device=device
|
|
).view(1,1,3,3)
|
|
|
|
sobel_y = torch.tensor(
|
|
[[-1,-2,-1],[0,0,0],[1,2,1]],
|
|
dtype=torch.float32, device=device
|
|
).view(1,1,3,3)
|
|
|
|
gx = F.conv2d(t, sobel_x, padding=1)
|
|
gy = F.conv2d(t, sobel_y, padding=1)
|
|
|
|
edge_energy = torch.mean(torch.sqrt(gx**2 + gy**2))
|
|
|
|
value = diff + 0.5 * edge_energy
|
|
|
|
signal_out.append(value.item())
|
|
|
|
prev_blocks = blocks
|
|
|
|
signal_out = np.array(signal_out)
|
|
|
|
# temporal derivative
|
|
signal_out = np.diff(signal_out, prepend=signal_out[0])
|
|
|
|
return signal_out
|
|
|
|
|
|
def normalize(sig):
|
|
return (sig - np.mean(sig)) / (np.std(sig) + 1e-6)
|
|
|
|
|
|
def find_offset(video_a, video_b):
|
|
|
|
sig_a = compute_signal(video_a)
|
|
sig_b = compute_signal(video_b)
|
|
|
|
sig_a = normalize(sig_a)
|
|
sig_b = normalize(sig_b)
|
|
|
|
corr = signal.correlate(sig_a, sig_b, mode="full")
|
|
lags = signal.correlation_lags(len(sig_a), len(sig_b), mode="full")
|
|
|
|
best = lags[np.argmax(corr)]
|
|
|
|
seconds = best / TARGET_FPS
|
|
|
|
return best, seconds
|
|
|
|
|
|
def main():
|
|
|
|
video_a = "/home/honney/projects/mkv/Maid-sama/11_org.mkv"
|
|
video_b = "/home/honney/projects/mkv/Maid-sama/11_ger.mkv"
|
|
|
|
frames, seconds = find_offset(video_a, video_b)
|
|
|
|
print("offset frames:", frames)
|
|
print("offset seconds:", seconds)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |