|
|
|
|
|
import os, pdb, subprocess, glob, cv2
|
|
import numpy as np
|
|
from shutil import rmtree
|
|
import torch
|
|
|
|
from scenedetect.video_manager import VideoManager
|
|
from scenedetect.scene_manager import SceneManager
|
|
from scenedetect.stats_manager import StatsManager
|
|
from scenedetect.detectors import ContentDetector
|
|
|
|
from scipy.interpolate import interp1d
|
|
from scipy.io import wavfile
|
|
from scipy import signal
|
|
|
|
from eval.detectors import S3FD
|
|
|
|
|
|
class SyncNetDetector:
|
|
def __init__(self, device, detect_results_dir="detect_results"):
|
|
self.s3f_detector = S3FD(device=device)
|
|
self.detect_results_dir = detect_results_dir
|
|
|
|
def __call__(self, video_path: str, min_track=50, scale=False):
|
|
crop_dir = os.path.join(self.detect_results_dir, "crop")
|
|
video_dir = os.path.join(self.detect_results_dir, "video")
|
|
frames_dir = os.path.join(self.detect_results_dir, "frames")
|
|
temp_dir = os.path.join(self.detect_results_dir, "temp")
|
|
|
|
|
|
if os.path.exists(crop_dir):
|
|
rmtree(crop_dir)
|
|
|
|
if os.path.exists(video_dir):
|
|
rmtree(video_dir)
|
|
|
|
if os.path.exists(frames_dir):
|
|
rmtree(frames_dir)
|
|
|
|
if os.path.exists(temp_dir):
|
|
rmtree(temp_dir)
|
|
|
|
|
|
|
|
os.makedirs(crop_dir)
|
|
os.makedirs(video_dir)
|
|
os.makedirs(frames_dir)
|
|
os.makedirs(temp_dir)
|
|
|
|
|
|
|
|
if scale:
|
|
scaled_video_path = os.path.join(video_dir, "scaled.mp4")
|
|
command = f"ffmpeg -loglevel error -y -nostdin -i {video_path} -vf scale='224:224' {scaled_video_path}"
|
|
subprocess.run(command, shell=True)
|
|
video_path = scaled_video_path
|
|
|
|
command = f"ffmpeg -y -nostdin -loglevel error -i {video_path} -qscale:v 2 -async 1 -r 25 {os.path.join(video_dir, 'video.mp4')}"
|
|
subprocess.run(command, shell=True, stdout=None)
|
|
|
|
command = f"ffmpeg -y -nostdin -loglevel error -i {os.path.join(video_dir, 'video.mp4')} -qscale:v 2 -f image2 {os.path.join(frames_dir, '%06d.jpg')}"
|
|
subprocess.run(command, shell=True, stdout=None)
|
|
|
|
command = f"ffmpeg -y -nostdin -loglevel error -i {os.path.join(video_dir, 'video.mp4')} -ac 1 -vn -acodec pcm_s16le -ar 16000 {os.path.join(video_dir, 'audio.wav')}"
|
|
subprocess.run(command, shell=True, stdout=None)
|
|
|
|
faces = self.detect_face(frames_dir)
|
|
|
|
scene = self.scene_detect(video_dir)
|
|
|
|
|
|
alltracks = []
|
|
|
|
for shot in scene:
|
|
if shot[1].frame_num - shot[0].frame_num >= min_track:
|
|
alltracks.extend(self.track_face(faces[shot[0].frame_num : shot[1].frame_num], min_track=min_track))
|
|
|
|
|
|
for ii, track in enumerate(alltracks):
|
|
self.crop_video(track, os.path.join(crop_dir, "%05d" % ii), frames_dir, 25, temp_dir, video_dir)
|
|
|
|
rmtree(temp_dir)
|
|
|
|
def scene_detect(self, video_dir):
|
|
video_manager = VideoManager([os.path.join(video_dir, "video.mp4")])
|
|
stats_manager = StatsManager()
|
|
scene_manager = SceneManager(stats_manager)
|
|
|
|
scene_manager.add_detector(ContentDetector())
|
|
base_timecode = video_manager.get_base_timecode()
|
|
|
|
video_manager.set_downscale_factor()
|
|
|
|
video_manager.start()
|
|
|
|
scene_manager.detect_scenes(frame_source=video_manager)
|
|
|
|
scene_list = scene_manager.get_scene_list(base_timecode)
|
|
|
|
if scene_list == []:
|
|
scene_list = [(video_manager.get_base_timecode(), video_manager.get_current_timecode())]
|
|
|
|
return scene_list
|
|
|
|
def track_face(self, scenefaces, num_failed_det=25, min_track=50, min_face_size=100):
|
|
|
|
iouThres = 0.5
|
|
tracks = []
|
|
|
|
while True:
|
|
track = []
|
|
for framefaces in scenefaces:
|
|
for face in framefaces:
|
|
if track == []:
|
|
track.append(face)
|
|
framefaces.remove(face)
|
|
elif face["frame"] - track[-1]["frame"] <= num_failed_det:
|
|
iou = bounding_box_iou(face["bbox"], track[-1]["bbox"])
|
|
if iou > iouThres:
|
|
track.append(face)
|
|
framefaces.remove(face)
|
|
continue
|
|
else:
|
|
break
|
|
|
|
if track == []:
|
|
break
|
|
elif len(track) > min_track:
|
|
|
|
framenum = np.array([f["frame"] for f in track])
|
|
bboxes = np.array([np.array(f["bbox"]) for f in track])
|
|
|
|
frame_i = np.arange(framenum[0], framenum[-1] + 1)
|
|
|
|
bboxes_i = []
|
|
for ij in range(0, 4):
|
|
interpfn = interp1d(framenum, bboxes[:, ij])
|
|
bboxes_i.append(interpfn(frame_i))
|
|
bboxes_i = np.stack(bboxes_i, axis=1)
|
|
|
|
if (
|
|
max(np.mean(bboxes_i[:, 2] - bboxes_i[:, 0]), np.mean(bboxes_i[:, 3] - bboxes_i[:, 1]))
|
|
> min_face_size
|
|
):
|
|
tracks.append({"frame": frame_i, "bbox": bboxes_i})
|
|
|
|
return tracks
|
|
|
|
def detect_face(self, frames_dir, facedet_scale=0.25):
|
|
flist = glob.glob(os.path.join(frames_dir, "*.jpg"))
|
|
flist.sort()
|
|
|
|
dets = []
|
|
|
|
for fidx, fname in enumerate(flist):
|
|
image = cv2.imread(fname)
|
|
|
|
image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
bboxes = self.s3f_detector.detect_faces(image_np, conf_th=0.9, scales=[facedet_scale])
|
|
|
|
dets.append([])
|
|
for bbox in bboxes:
|
|
dets[-1].append({"frame": fidx, "bbox": (bbox[:-1]).tolist(), "conf": bbox[-1]})
|
|
|
|
return dets
|
|
|
|
def crop_video(self, track, cropfile, frames_dir, frame_rate, temp_dir, video_dir, crop_scale=0.4):
|
|
|
|
flist = glob.glob(os.path.join(frames_dir, "*.jpg"))
|
|
flist.sort()
|
|
|
|
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
|
vOut = cv2.VideoWriter(cropfile + "t.mp4", fourcc, frame_rate, (224, 224))
|
|
|
|
dets = {"x": [], "y": [], "s": []}
|
|
|
|
for det in track["bbox"]:
|
|
|
|
dets["s"].append(max((det[3] - det[1]), (det[2] - det[0])) / 2)
|
|
dets["y"].append((det[1] + det[3]) / 2)
|
|
dets["x"].append((det[0] + det[2]) / 2)
|
|
|
|
|
|
dets["s"] = signal.medfilt(dets["s"], kernel_size=13)
|
|
dets["x"] = signal.medfilt(dets["x"], kernel_size=13)
|
|
dets["y"] = signal.medfilt(dets["y"], kernel_size=13)
|
|
|
|
for fidx, frame in enumerate(track["frame"]):
|
|
|
|
cs = crop_scale
|
|
|
|
bs = dets["s"][fidx]
|
|
bsi = int(bs * (1 + 2 * cs))
|
|
|
|
image = cv2.imread(flist[frame])
|
|
|
|
frame = np.pad(image, ((bsi, bsi), (bsi, bsi), (0, 0)), "constant", constant_values=(110, 110))
|
|
my = dets["y"][fidx] + bsi
|
|
mx = dets["x"][fidx] + bsi
|
|
|
|
face = frame[int(my - bs) : int(my + bs * (1 + 2 * cs)), int(mx - bs * (1 + cs)) : int(mx + bs * (1 + cs))]
|
|
|
|
vOut.write(cv2.resize(face, (224, 224)))
|
|
|
|
audiotmp = os.path.join(temp_dir, "audio.wav")
|
|
audiostart = (track["frame"][0]) / frame_rate
|
|
audioend = (track["frame"][-1] + 1) / frame_rate
|
|
|
|
vOut.release()
|
|
|
|
|
|
|
|
command = "ffmpeg -y -nostdin -loglevel error -i %s -ss %.3f -to %.3f %s" % (
|
|
os.path.join(video_dir, "audio.wav"),
|
|
audiostart,
|
|
audioend,
|
|
audiotmp,
|
|
)
|
|
output = subprocess.run(command, shell=True, stdout=None)
|
|
|
|
sample_rate, audio = wavfile.read(audiotmp)
|
|
|
|
|
|
|
|
command = "ffmpeg -y -nostdin -loglevel error -i %st.mp4 -i %s -c:v copy -c:a aac %s.mp4" % (
|
|
cropfile,
|
|
audiotmp,
|
|
cropfile,
|
|
)
|
|
output = subprocess.run(command, shell=True, stdout=None)
|
|
|
|
os.remove(cropfile + "t.mp4")
|
|
|
|
return {"track": track, "proc_track": dets}
|
|
|
|
|
|
def bounding_box_iou(boxA, boxB):
|
|
xA = max(boxA[0], boxB[0])
|
|
yA = max(boxA[1], boxB[1])
|
|
xB = min(boxA[2], boxB[2])
|
|
yB = min(boxA[3], boxB[3])
|
|
|
|
interArea = max(0, xB - xA) * max(0, yB - yA)
|
|
|
|
boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
|
|
boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
|
|
|
|
iou = interArea / float(boxAArea + boxBArea - interArea)
|
|
|
|
return iou
|
|
|