svjack/LatentSync
0
1# Adapted from https://github.com/joonson/syncnet_python/blob/master/run_pipeline.py2 3import os, pdb, subprocess, glob, cv24import numpy as np5from shutil import rmtree6import torch7 8from scenedetect.video_manager import VideoManager9from scenedetect.scene_manager import SceneManager10from scenedetect.stats_manager import StatsManager11from scenedetect.detectors import ContentDetector12 13from scipy.interpolate import interp1d14from scipy.io import wavfile15from scipy import signal16 17from eval.detectors import S3FD18 19 20class SyncNetDetector:21 def __init__(self, device, detect_results_dir="detect_results"):22 self.s3f_detector = S3FD(device=device)23 self.detect_results_dir = detect_results_dir24 25 def __call__(self, video_path: str, min_track=50, scale=False):26 crop_dir = os.path.join(self.detect_results_dir, "crop")27 video_dir = os.path.join(self.detect_results_dir, "video")28 frames_dir = os.path.join(self.detect_results_dir, "frames")29 temp_dir = os.path.join(self.detect_results_dir, "temp")30 31 # ========== DELETE EXISTING DIRECTORIES ==========32 if os.path.exists(crop_dir):33 rmtree(crop_dir)34 35 if os.path.exists(video_dir):36 rmtree(video_dir)37 38 if os.path.exists(frames_dir):39 rmtree(frames_dir)40 41 if os.path.exists(temp_dir):42 rmtree(temp_dir)43 44 # ========== MAKE NEW DIRECTORIES ==========45 46 os.makedirs(crop_dir)47 os.makedirs(video_dir)48 os.makedirs(frames_dir)49 os.makedirs(temp_dir)50 51 # ========== CONVERT VIDEO AND EXTRACT FRAMES ==========52 53 if scale:54 scaled_video_path = os.path.join(video_dir, "scaled.mp4")55 command = f"ffmpeg -loglevel error -y -nostdin -i {video_path} -vf scale='224:224' {scaled_video_path}"56 subprocess.run(command, shell=True)57 video_path = scaled_video_path58 59 command = f"ffmpeg -y -nostdin -loglevel error -i {video_path} -qscale:v 2 -async 1 -r 25 {os.path.join(video_dir, 'video.mp4')}"60 subprocess.run(command, shell=True, stdout=None)61 62 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')}"63 subprocess.run(command, shell=True, stdout=None)64 65 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')}"66 subprocess.run(command, shell=True, stdout=None)67 68 faces = self.detect_face(frames_dir)69 70 scene = self.scene_detect(video_dir)71 72 # Face tracking73 alltracks = []74 75 for shot in scene:76 if shot[1].frame_num - shot[0].frame_num >= min_track:77 alltracks.extend(self.track_face(faces[shot[0].frame_num : shot[1].frame_num], min_track=min_track))78 79 # Face crop80 for ii, track in enumerate(alltracks):81 self.crop_video(track, os.path.join(crop_dir, "%05d" % ii), frames_dir, 25, temp_dir, video_dir)82 83 rmtree(temp_dir)84 85 def scene_detect(self, video_dir):86 video_manager = VideoManager([os.path.join(video_dir, "video.mp4")])87 stats_manager = StatsManager()88 scene_manager = SceneManager(stats_manager)89 # Add ContentDetector algorithm (constructor takes detector options like threshold).90 scene_manager.add_detector(ContentDetector())91 base_timecode = video_manager.get_base_timecode()92 93 video_manager.set_downscale_factor()94 95 video_manager.start()96 97 scene_manager.detect_scenes(frame_source=video_manager)98 99 scene_list = scene_manager.get_scene_list(base_timecode)100 101 if scene_list == []:102 scene_list = [(video_manager.get_base_timecode(), video_manager.get_current_timecode())]103 104 return scene_list105 106 def track_face(self, scenefaces, num_failed_det=25, min_track=50, min_face_size=100):107 108 iouThres = 0.5 # Minimum IOU between consecutive face detections109 tracks = []110 111 while True:112 track = []113 for framefaces in scenefaces:114 for face in framefaces:115 if track == []:116 track.append(face)117 framefaces.remove(face)118 elif face["frame"] - track[-1]["frame"] <= num_failed_det:119 iou = bounding_box_iou(face["bbox"], track[-1]["bbox"])120 if iou > iouThres:121 track.append(face)122 framefaces.remove(face)123 continue124 else:125 break126 127 if track == []:128 break129 elif len(track) > min_track:130 131 framenum = np.array([f["frame"] for f in track])132 bboxes = np.array([np.array(f["bbox"]) for f in track])133 134 frame_i = np.arange(framenum[0], framenum[-1] + 1)135 136 bboxes_i = []137 for ij in range(0, 4):138 interpfn = interp1d(framenum, bboxes[:, ij])139 bboxes_i.append(interpfn(frame_i))140 bboxes_i = np.stack(bboxes_i, axis=1)141 142 if (143 max(np.mean(bboxes_i[:, 2] - bboxes_i[:, 0]), np.mean(bboxes_i[:, 3] - bboxes_i[:, 1]))144 > min_face_size145 ):146 tracks.append({"frame": frame_i, "bbox": bboxes_i})147 148 return tracks149 150 def detect_face(self, frames_dir, facedet_scale=0.25):151 flist = glob.glob(os.path.join(frames_dir, "*.jpg"))152 flist.sort()153 154 dets = []155 156 for fidx, fname in enumerate(flist):157 image = cv2.imread(fname)158 159 image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)160 bboxes = self.s3f_detector.detect_faces(image_np, conf_th=0.9, scales=[facedet_scale])161 162 dets.append([])163 for bbox in bboxes:164 dets[-1].append({"frame": fidx, "bbox": (bbox[:-1]).tolist(), "conf": bbox[-1]})165 166 return dets167 168 def crop_video(self, track, cropfile, frames_dir, frame_rate, temp_dir, video_dir, crop_scale=0.4):169 170 flist = glob.glob(os.path.join(frames_dir, "*.jpg"))171 flist.sort()172 173 fourcc = cv2.VideoWriter_fourcc(*"mp4v")174 vOut = cv2.VideoWriter(cropfile + "t.mp4", fourcc, frame_rate, (224, 224))175 176 dets = {"x": [], "y": [], "s": []}177 178 for det in track["bbox"]:179 180 dets["s"].append(max((det[3] - det[1]), (det[2] - det[0])) / 2)181 dets["y"].append((det[1] + det[3]) / 2) # crop center x182 dets["x"].append((det[0] + det[2]) / 2) # crop center y183 184 # Smooth detections185 dets["s"] = signal.medfilt(dets["s"], kernel_size=13)186 dets["x"] = signal.medfilt(dets["x"], kernel_size=13)187 dets["y"] = signal.medfilt(dets["y"], kernel_size=13)188 189 for fidx, frame in enumerate(track["frame"]):190 191 cs = crop_scale192 193 bs = dets["s"][fidx] # Detection box size194 bsi = int(bs * (1 + 2 * cs)) # Pad videos by this amount195 196 image = cv2.imread(flist[frame])197 198 frame = np.pad(image, ((bsi, bsi), (bsi, bsi), (0, 0)), "constant", constant_values=(110, 110))199 my = dets["y"][fidx] + bsi # BBox center Y200 mx = dets["x"][fidx] + bsi # BBox center X201 202 face = frame[int(my - bs) : int(my + bs * (1 + 2 * cs)), int(mx - bs * (1 + cs)) : int(mx + bs * (1 + cs))]203 204 vOut.write(cv2.resize(face, (224, 224)))205 206 audiotmp = os.path.join(temp_dir, "audio.wav")207 audiostart = (track["frame"][0]) / frame_rate208 audioend = (track["frame"][-1] + 1) / frame_rate209 210 vOut.release()211 212 # ========== CROP AUDIO FILE ==========213 214 command = "ffmpeg -y -nostdin -loglevel error -i %s -ss %.3f -to %.3f %s" % (215 os.path.join(video_dir, "audio.wav"),216 audiostart,217 audioend,218 audiotmp,219 )220 output = subprocess.run(command, shell=True, stdout=None)221 222 sample_rate, audio = wavfile.read(audiotmp)223 224 # ========== COMBINE AUDIO AND VIDEO FILES ==========225 226 command = "ffmpeg -y -nostdin -loglevel error -i %st.mp4 -i %s -c:v copy -c:a aac %s.mp4" % (227 cropfile,228 audiotmp,229 cropfile,230 )231 output = subprocess.run(command, shell=True, stdout=None)232 233 os.remove(cropfile + "t.mp4")234 235 return {"track": track, "proc_track": dets}236 237 238def bounding_box_iou(boxA, boxB):239 xA = max(boxA[0], boxB[0])240 yA = max(boxA[1], boxB[1])241 xB = min(boxA[2], boxB[2])242 yB = min(boxA[3], boxB[3])243 244 interArea = max(0, xB - xA) * max(0, yB - yA)245 246 boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])247 boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])248 249 iou = interArea / float(boxAArea + boxBArea - interArea)250 251 return iou252