svjack/LatentSync

1# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import argparse16import os17import tqdm18from statistics import fmean19from eval.syncnet import SyncNetEval20from eval.syncnet_detect import SyncNetDetector21from latentsync.utils.util import red_text22import torch23 24 25def syncnet_eval(syncnet, syncnet_detector, video_path, temp_dir, detect_results_dir="detect_results"):26    syncnet_detector(video_path=video_path, min_track=50)27    crop_videos = os.listdir(os.path.join(detect_results_dir, "crop"))28    if crop_videos == []:29        raise Exception(red_text(f"Face not detected in {video_path}"))30    av_offset_list = []31    conf_list = []32    for video in crop_videos:33        av_offset, _, conf = syncnet.evaluate(34            video_path=os.path.join(detect_results_dir, "crop", video), temp_dir=temp_dir35        )36        av_offset_list.append(av_offset)37        conf_list.append(conf)38    av_offset = int(fmean(av_offset_list))39    conf = fmean(conf_list)40    print(f"Input video: {video_path}\nSyncNet confidence: {conf:.2f}\nAV offset: {av_offset}")41    return av_offset, conf42 43 44def main():45    parser = argparse.ArgumentParser(description="SyncNet")46    parser.add_argument("--initial_model", type=str, default="checkpoints/auxiliary/syncnet_v2.model", help="")47    parser.add_argument("--video_path", type=str, default=None, help="")48    parser.add_argument("--videos_dir", type=str, default="/root/processed")49    parser.add_argument("--temp_dir", type=str, default="temp", help="")50 51    args = parser.parse_args()52 53    device = "cuda" if torch.cuda.is_available() else "cpu"54 55    syncnet = SyncNetEval(device=device)56    syncnet.loadParameters(args.initial_model)57 58    syncnet_detector = SyncNetDetector(device=device, detect_results_dir="detect_results")59 60    if args.video_path is not None:61        syncnet_eval(syncnet, syncnet_detector, args.video_path, args.temp_dir)62    else:63        sync_conf_list = []64        video_names = sorted([f for f in os.listdir(args.videos_dir) if f.endswith(".mp4")])65        for video_name in tqdm.tqdm(video_names):66            try:67                _, conf = syncnet_eval(68                    syncnet, syncnet_detector, os.path.join(args.videos_dir, video_name), args.temp_dir69                )70                sync_conf_list.append(conf)71            except Exception as e:72                print(e)73        print(f"The average sync confidence is {fmean(sync_conf_list):.02f}")74 75 76if __name__ == "__main__":77    main()78