pandora/pandora/archive/cutdetection.py
2026-01-15 15:59:14 +00:00

109 lines
3.4 KiB
Python

import subprocess
import numpy as np
import ox
from django.conf import settings
from .utils import AspectRatio
def _get_distance(data0, data1):
diff = data0.astype(np.float32) - data1.astype(np.float32)
per_pixel_distance = np.linalg.norm(diff, axis=2)
total_distance = per_pixel_distance.sum()
num_pixels = data0.shape[0] * data0.shape[1]
max_distance = num_pixels * np.sqrt(3 * (255 ** 2))
return total_distance / max_distance
def detect_cuts(path, seconds=True):
if isinstance(path, list):
cuts = []
position = 0
for segment in path:
info = ox.avinfo(segment)
segment_cuts = detect_cuts(segment, seconds)
if segment_cuts:
if seconds:
if cuts:
cuts.append(position)
cuts += [c + position for c in segment_cuts]
else:
if cuts:
cuts.append(int(position*fps))
fps = AspectRatio(info['video'][0]['framerate'])
cuts += [c + int(position*fps) for c in segment_cuts]
position += info['duration']
return cuts
depth = 3
info = ox.avinfo(path)
if not info.get('video'):
return []
dar = AspectRatio(info['video'][0]['display_aspect_ratio'])
fps = AspectRatio(info['video'][0]['framerate'])
height = 96
width = int(dar * height)
width += width % 2
nbytes = depth * width * height
bufsize = nbytes + 100
cmd = [
settings.FFMPEG,
'-hide_banner',
'-loglevel', 'error',
'-i', path,
'-threads', '4',
'-f', 'rawvideo',
'-pix_fmt', 'rgb24',
'-vcodec', 'rawvideo',
'-vf', 'scale=%d:%d' % (width, height),
'-aspect', '%d:%d' % (width, height),
'-'
]
#print(' '.join(cmd))
p = subprocess.Popen(cmd,
bufsize=bufsize,
stdout=subprocess.PIPE,
close_fds=True)
first = True
previous_frame = None
cuts = []
short_cut = None
cut_frames = []
frame_i = 0
previous_distance = 0
while True:
data = p.stdout.read(nbytes)
if len(data) != nbytes:
if first:
raise IOError("ERROR: could not open file %s" % path)
else:
break
else:
first = False
frame = np.frombuffer(data, dtype=np.uint8).reshape((height, width, depth))
frame_data = np.asarray(frame)
if frame_i == 0:
is_cut = False
previous_distance = 0
else:
if short_cut and frame_i - short_cut > 2:
cuts.append(short_cut)
short_cut = None
distance = _get_distance(previous_frame_data, frame_data)
is_cut = distance > 0.1 or (distance > 0.2 and abs(distance - previous_distance) > 0.2)
if is_cut:
if frame_i - (0 if not cuts else short_cut or cuts[-1]) < 3:
is_cut = False
short_cut = frame_i
else:
cuts.append(frame_i)
previous_distance = distance
previous_frame_data = frame_data
frame_i += 1
if seconds:
return [float('%0.3f' % float(c/fps)) for c in cuts]
else:
return cuts