# Adapted from https://github.com/joonson/syncnet_python/blob/master/run_pipeline.py 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") # ========== DELETE EXISTING DIRECTORIES ========== 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) # ========== MAKE NEW DIRECTORIES ========== os.makedirs(crop_dir) os.makedirs(video_dir) os.makedirs(frames_dir) os.makedirs(temp_dir) # ========== CONVERT VIDEO AND EXTRACT FRAMES ========== 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) # Face tracking 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)) # Face crop 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) # Add ContentDetector algorithm (constructor takes detector options like threshold). 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 # Minimum IOU between consecutive face detections 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) # crop center x dets["x"].append((det[0] + det[2]) / 2) # crop center y # Smooth detections 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] # Detection box size bsi = int(bs * (1 + 2 * cs)) # Pad videos by this amount 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 # BBox center Y mx = dets["x"][fidx] + bsi # BBox center X 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() # ========== CROP AUDIO FILE ========== 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) # ========== COMBINE AUDIO AND VIDEO FILES ========== 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