From 4050b05a4071a322758b1546f2a06e64e8f8a4d5 Mon Sep 17 00:00:00 2001 From: zodymm Date: Fri, 9 May 2025 11:36:54 +0800 Subject: [PATCH 1/3] update multiprocessing --- SyncNetInstance.py | 207 +++++++++++++---------- crop_runner.py | 414 +++++++++++++++++++++++++++++++++++++++++++++ syncnet_runner.py | 163 ++++++++++++++++++ 3 files changed, 691 insertions(+), 93 deletions(-) create mode 100755 crop_runner.py create mode 100755 syncnet_runner.py diff --git a/SyncNetInstance.py b/SyncNetInstance.py index 497d44f..2afb111 100644 --- a/SyncNetInstance.py +++ b/SyncNetInstance.py @@ -1,78 +1,98 @@ #!/usr/bin/python -#-*- coding: utf-8 -*- +# -*- coding: utf-8 -*- # Video 25 FPS, Audio 16000HZ -import torch -import numpy -import time, pdb, argparse, subprocess, os, math, glob +import argparse +import glob +import math +import os +import pdb +import subprocess +import time +from shutil import rmtree + import cv2 +import numpy import python_speech_features - +import torch from scipy import signal from scipy.io import wavfile + from SyncNetModel import * -from shutil import rmtree + + +def silent_call(cmd): + return subprocess.call(cmd, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) # ==================== Get OFFSET ==================== + def calc_pdist(feat1, feat2, vshift=10): - - win_size = vshift*2+1 + win_size = vshift * 2 + 1 - feat2p = torch.nn.functional.pad(feat2,(0,0,vshift,vshift)) + feat2p = torch.nn.functional.pad(feat2, (0, 0, vshift, vshift)) dists = [] - for i in range(0,len(feat1)): - - dists.append(torch.nn.functional.pairwise_distance(feat1[[i],:].repeat(win_size, 1), feat2p[i:i+win_size,:])) + for i in range(0, len(feat1)): + dists.append( + torch.nn.functional.pairwise_distance(feat1[[i], :].repeat(win_size, 1), feat2p[i : i + win_size, :]) + ) return dists + # ==================== MAIN DEF ==================== -class SyncNetInstance(torch.nn.Module): - def __init__(self, dropout = 0, num_layers_in_fc_layers = 1024): - super(SyncNetInstance, self).__init__(); +class SyncNetInstance(torch.nn.Module): + def __init__(self, dropout=0, num_layers_in_fc_layers=1024, device="cpu"): + super(SyncNetInstance, self).__init__() + self.__S__ = S(num_layers_in_fc_layers=num_layers_in_fc_layers).to(device) - self.__S__ = S(num_layers_in_fc_layers = num_layers_in_fc_layers).cuda(); + self.device = device def evaluate(self, opt, videofile): - - self.__S__.eval(); - + self.__S__.eval() # ========== ========== # Convert files # ========== ========== - if os.path.exists(os.path.join(opt.tmp_dir,opt.reference)): - rmtree(os.path.join(opt.tmp_dir,opt.reference)) + if os.path.exists(os.path.join(opt.tmp_dir, opt.reference)): + rmtree(os.path.join(opt.tmp_dir, opt.reference)) + + os.makedirs(os.path.join(opt.tmp_dir, opt.reference)) - os.makedirs(os.path.join(opt.tmp_dir,opt.reference)) + command = "ffmpeg -y -i %s -threads 1 -f image2 %s" % ( + videofile, + os.path.join(opt.tmp_dir, opt.reference, "%06d.jpg"), + ) + # output = subprocess.call(command, shell=True, stdout=None) + silent_call(command) - command = ("ffmpeg -y -i %s -threads 1 -f image2 %s" % (videofile,os.path.join(opt.tmp_dir,opt.reference,'%06d.jpg'))) - output = subprocess.call(command, shell=True, stdout=None) + command = "ffmpeg -y -i %s -async 1 -ac 1 -vn -acodec pcm_s16le -ar 16000 %s" % ( + videofile, + os.path.join(opt.tmp_dir, opt.reference, "audio.wav"), + ) + # output = subprocess.call(command, shell=True, stdout=None) + silent_call(command) - command = ("ffmpeg -y -i %s -async 1 -ac 1 -vn -acodec pcm_s16le -ar 16000 %s" % (videofile,os.path.join(opt.tmp_dir,opt.reference,'audio.wav'))) - output = subprocess.call(command, shell=True, stdout=None) - # ========== ========== - # Load video + # Load video # ========== ========== images = [] - - flist = glob.glob(os.path.join(opt.tmp_dir,opt.reference,'*.jpg')) + + flist = glob.glob(os.path.join(opt.tmp_dir, opt.reference, "*.jpg")) flist.sort() for fname in flist: images.append(cv2.imread(fname)) - im = numpy.stack(images,axis=3) - im = numpy.expand_dims(im,axis=0) - im = numpy.transpose(im,(0,3,4,1,2)) + im = numpy.stack(images, axis=3) + im = numpy.expand_dims(im, axis=0) + im = numpy.transpose(im, (0, 3, 4, 1, 2)) imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float()) @@ -80,83 +100,87 @@ def evaluate(self, opt, videofile): # Load audio # ========== ========== - sample_rate, audio = wavfile.read(os.path.join(opt.tmp_dir,opt.reference,'audio.wav')) - mfcc = zip(*python_speech_features.mfcc(audio,sample_rate)) + sample_rate, audio = wavfile.read(os.path.join(opt.tmp_dir, opt.reference, "audio.wav")) + mfcc = zip(*python_speech_features.mfcc(audio, sample_rate)) mfcc = numpy.stack([numpy.array(i) for i in mfcc]) - cc = numpy.expand_dims(numpy.expand_dims(mfcc,axis=0),axis=0) + cc = numpy.expand_dims(numpy.expand_dims(mfcc, axis=0), axis=0) cct = torch.autograd.Variable(torch.from_numpy(cc.astype(float)).float()) # ========== ========== # Check audio and video input length # ========== ========== - if (float(len(audio))/16000) != (float(len(images))/25) : - print("WARNING: Audio (%.4fs) and video (%.4fs) lengths are different."%(float(len(audio))/16000,float(len(images))/25)) + # if (float(len(audio)) / 16000) != (float(len(images)) / 25): + # print( + # "WARNING: Audio (%.4fs) and video (%.4fs) lengths are different." + # % (float(len(audio)) / 16000, float(len(images)) / 25) + # ) + + min_length = min(len(images), math.floor(len(audio) / 640)) - min_length = min(len(images),math.floor(len(audio)/640)) - # ========== ========== # Generate video and audio feats # ========== ========== - lastframe = min_length-5 + lastframe = min_length - 5 im_feat = [] cc_feat = [] tS = time.time() - for i in range(0,lastframe,opt.batch_size): - - im_batch = [ imtv[:,:,vframe:vframe+5,:,:] for vframe in range(i,min(lastframe,i+opt.batch_size)) ] - im_in = torch.cat(im_batch,0) - im_out = self.__S__.forward_lip(im_in.cuda()); + for i in range(0, lastframe, opt.batch_size): + im_batch = [ + imtv[:, :, vframe : vframe + 5, :, :] for vframe in range(i, min(lastframe, i + opt.batch_size)) + ] + im_in = torch.cat(im_batch, 0) + im_out = self.__S__.forward_lip(im_in.to(self.device)) im_feat.append(im_out.data.cpu()) - cc_batch = [ cct[:,:,:,vframe*4:vframe*4+20] for vframe in range(i,min(lastframe,i+opt.batch_size)) ] - cc_in = torch.cat(cc_batch,0) - cc_out = self.__S__.forward_aud(cc_in.cuda()) + cc_batch = [ + cct[:, :, :, vframe * 4 : vframe * 4 + 20] for vframe in range(i, min(lastframe, i + opt.batch_size)) + ] + cc_in = torch.cat(cc_batch, 0) + cc_out = self.__S__.forward_aud(cc_in.to(self.device)) cc_feat.append(cc_out.data.cpu()) - im_feat = torch.cat(im_feat,0) - cc_feat = torch.cat(cc_feat,0) + im_feat = torch.cat(im_feat, 0) + cc_feat = torch.cat(cc_feat, 0) # ========== ========== # Compute offset # ========== ========== - - print('Compute time %.3f sec.' % (time.time()-tS)) - dists = calc_pdist(im_feat,cc_feat,vshift=opt.vshift) - mdist = torch.mean(torch.stack(dists,1),1) + # print("Compute time %.3f sec." % (time.time() - tS)) - minval, minidx = torch.min(mdist,0) + dists = calc_pdist(im_feat, cc_feat, vshift=opt.vshift) + mdist = torch.mean(torch.stack(dists, 1), 1) - offset = opt.vshift-minidx - conf = torch.median(mdist) - minval + minval, minidx = torch.min(mdist, 0) - fdist = numpy.stack([dist[minidx].numpy() for dist in dists]) + offset = opt.vshift - minidx + conf = torch.median(mdist) - minval + + fdist = numpy.stack([dist[minidx].numpy() for dist in dists]) # fdist = numpy.pad(fdist, (3,3), 'constant', constant_values=15) - fconf = torch.median(mdist).numpy() - fdist - fconfm = signal.medfilt(fconf,kernel_size=9) - - numpy.set_printoptions(formatter={'float': '{: 0.3f}'.format}) - print('Framewise conf: ') - print(fconfm) - print('AV offset: \t%d \nMin dist: \t%.3f\nConfidence: \t%.3f' % (offset,minval,conf)) + fconf = torch.median(mdist).numpy() - fdist + fconfm = signal.medfilt(fconf, kernel_size=9) - dists_npy = numpy.array([ dist.numpy() for dist in dists ]) - return offset.numpy(), conf.numpy(), dists_npy + numpy.set_printoptions(formatter={"float": "{: 0.3f}".format}) + # print("Framewise conf: ") + # print(fconfm) + # print("AV offset: \t%d \nMin dist: \t%.3f\nConfidence: \t%.3f" % (offset, minval, conf)) - def extract_feature(self, opt, videofile): + dists_npy = numpy.array([dist.numpy() for dist in dists]) + return offset.numpy(), conf.numpy(), dists_npy, minval.numpy() - self.__S__.eval(); - + def extract_feature(self, opt, videofile): + self.__S__.eval() # ========== ========== - # Load video + # Load video # ========== ========== cap = cv2.VideoCapture(videofile) - frame_num = 1; + frame_num = 1 images = [] while frame_num: frame_num += 1 @@ -166,43 +190,40 @@ def extract_feature(self, opt, videofile): images.append(image) - im = numpy.stack(images,axis=3) - im = numpy.expand_dims(im,axis=0) - im = numpy.transpose(im,(0,3,4,1,2)) + im = numpy.stack(images, axis=3) + im = numpy.expand_dims(im, axis=0) + im = numpy.transpose(im, (0, 3, 4, 1, 2)) imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float()) - + # ========== ========== # Generate video feats # ========== ========== - lastframe = len(images)-4 + lastframe = len(images) - 4 im_feat = [] tS = time.time() - for i in range(0,lastframe,opt.batch_size): - - im_batch = [ imtv[:,:,vframe:vframe+5,:,:] for vframe in range(i,min(lastframe,i+opt.batch_size)) ] - im_in = torch.cat(im_batch,0) - im_out = self.__S__.forward_lipfeat(im_in.cuda()); + for i in range(0, lastframe, opt.batch_size): + im_batch = [ + imtv[:, :, vframe : vframe + 5, :, :] for vframe in range(i, min(lastframe, i + opt.batch_size)) + ] + im_in = torch.cat(im_batch, 0) + im_out = self.__S__.forward_lipfeat(im_in.cuda()) im_feat.append(im_out.data.cpu()) - im_feat = torch.cat(im_feat,0) + im_feat = torch.cat(im_feat, 0) # ========== ========== # Compute offset # ========== ========== - - print('Compute time %.3f sec.' % (time.time()-tS)) - return im_feat + # print("Compute time %.3f sec." % (time.time() - tS)) + return im_feat def loadParameters(self, path): - loaded_state = torch.load(path, map_location=lambda storage, loc: storage); - - self_state = self.__S__.state_dict(); - + loaded_state = torch.load(path, map_location=lambda storage, loc: storage) + self_state = self.__S__.state_dict() for name, param in loaded_state.items(): - - self_state[name].copy_(param); + self_state[name].copy_(param) diff --git a/crop_runner.py b/crop_runner.py new file mode 100755 index 0000000..4718a8c --- /dev/null +++ b/crop_runner.py @@ -0,0 +1,414 @@ +#!/usr/bin/python +import warnings + +warnings.filterwarnings("ignore") + +import argparse +import glob +import os +import pickle +import subprocess +import time +from pathlib import Path +from shutil import rmtree + +import cv2 +import numpy as np +import torch.multiprocessing as mp +from scenedetect.detectors import ContentDetector +from scenedetect.scene_manager import SceneManager +from scenedetect.stats_manager import StatsManager +from scenedetect.video_manager import VideoManager +from scipy import signal +from scipy.interpolate import interp1d +from scipy.io import wavfile +from tqdm import tqdm + +from detectors import S3FD + +# ========== ========== ========== ========== +# # IOU FUNCTION +# ========== ========== ========== ========== + + +def bb_intersection_over_union(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 + + +# ========== ========== ========== ========== +# # FACE TRACKING +# ========== ========== ========== ========== + + +def track_shot(opt, scenefaces): + 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"] <= opt.num_failed_det: + iou = bb_intersection_over_union(face["bbox"], track[-1]["bbox"]) + if iou > iouThres: + track.append(face) + framefaces.remove(face) + continue + else: + break + + if track == []: + break + elif len(track) > opt.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])) + > opt.min_face_size + ): + tracks.append({"frame": frame_i, "bbox": bboxes_i}) + + return tracks + + +# ========== ========== ========== ========== +# # VIDEO CROP AND SAVE +# ========== ========== ========== ========== + + +def crop_video(opt, track, cropfile): + flist = glob.glob(os.path.join(opt.frames_dir, opt.reference, "*.jpg")) + flist.sort() + + fourcc = cv2.VideoWriter_fourcc(*"XVID") + vOut = cv2.VideoWriter(cropfile + "t.avi", fourcc, opt.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 = opt.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(opt.tmp_dir, opt.reference, "audio.wav") + audiostart = (track["frame"][0]) / opt.frame_rate + audioend = (track["frame"][-1] + 1) / opt.frame_rate + + vOut.release() + + # ========== CROP AUDIO FILE ========== + + command = "ffmpeg -y -i %s -ss %.3f -to %.3f %s" % ( + os.path.join(opt.avi_dir, opt.reference, "audio.wav"), + audiostart, + audioend, + audiotmp, + ) + # output = subprocess.call(command, shell=True, stdout=None) + silent_call(command) + + # if output != 0: + # pdb.set_trace() + + sample_rate, audio = wavfile.read(audiotmp) + + # ========== COMBINE AUDIO AND VIDEO FILES ========== + + command = "ffmpeg -y -i %st.avi -i %s -c:v copy -c:a copy %s.avi" % (cropfile, audiotmp, cropfile) + # output = subprocess.call(command, shell=True, stdout=None) + silent_call(command) + + # if output != 0: + # pdb.set_trace() + + # print("Written %s" % cropfile) + + os.remove(cropfile + "t.avi") + + # print("Mean pos: x %.2f y %.2f s %.2f" % (np.mean(dets["x"]), np.mean(dets["y"]), np.mean(dets["s"]))) + + return {"track": track, "proc_track": dets} + + +# ========== ========== ========== ========== +# # FACE DETECTION +# ========== ========== ========== ========== + + +def inference_video(opt, det_model): + flist = glob.glob(os.path.join(opt.frames_dir, opt.reference, "*.jpg")) + flist.sort() + + dets = [] + + for fidx, fname in enumerate(flist): + start_time = time.time() + + image = cv2.imread(fname) + + image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + bboxes = det_model.detect_faces(image_np, conf_th=0.9, scales=[opt.facedet_scale]) + + dets.append([]) + for bbox in bboxes: + dets[-1].append({"frame": fidx, "bbox": (bbox[:-1]).tolist(), "conf": bbox[-1]}) + + elapsed_time = time.time() - start_time + + # print( + # "%s-%05d; %d dets; %.2f Hz" + # % (os.path.join(opt.avi_dir, opt.reference, "video.avi"), fidx, len(dets[-1]), (1 / elapsed_time)) + # ) + + savepath = os.path.join(opt.work_dir, opt.reference, "faces.pckl") + + with open(savepath, "wb") as fil: + pickle.dump(dets, fil) + + return dets + + +# ========== ========== ========== ========== +# # SCENE DETECTION +# ========== ========== ========== ========== + + +def scene_detect(opt): + video_manager = VideoManager([os.path.join(opt.avi_dir, opt.reference, "video.avi")]) + 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, show_progress=False) + + scene_list = scene_manager.get_scene_list(base_timecode) + + savepath = os.path.join(opt.work_dir, opt.reference, "scene.pckl") + + if scene_list == []: + scene_list = [(video_manager.get_base_timecode(), video_manager.get_current_timecode())] + + with open(savepath, "wb") as fil: + pickle.dump(scene_list, fil) + + # print("%s - scenes detected %d" % (os.path.join(opt.avi_dir, opt.reference, "video.avi"), len(scene_list))) + + return scene_list + + +def silent_call(cmd): + return subprocess.call(cmd, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + + +def parse_args(data_dir, videofile, reference): + opt = argparse.Namespace() + setattr(opt, "data_dir", data_dir) + setattr(opt, "videofile", videofile) + setattr(opt, "reference", reference) + setattr(opt, "facedet_scale", 0.25) + setattr(opt, "crop_scale", 0.40) + setattr(opt, "min_track", 30) + setattr(opt, "frame_rate", 25) + setattr(opt, "num_failed_det", 25) + setattr(opt, "min_face_size", 100) + + setattr(opt, "avi_dir", os.path.join(opt.data_dir, "pyavi")) + setattr(opt, "tmp_dir", os.path.join(opt.data_dir, "pytmp")) + setattr(opt, "work_dir", os.path.join(opt.data_dir, "pywork")) + setattr(opt, "crop_dir", os.path.join(opt.data_dir, "pycrop")) + setattr(opt, "frames_dir", os.path.join(opt.data_dir, "pyframes")) + + # ========== DELETE EXISTING DIRECTORIES ========== + + if os.path.exists(os.path.join(opt.work_dir, opt.reference)): + rmtree(os.path.join(opt.work_dir, opt.reference)) + + if os.path.exists(os.path.join(opt.crop_dir, opt.reference)): + rmtree(os.path.join(opt.crop_dir, opt.reference)) + + if os.path.exists(os.path.join(opt.avi_dir, opt.reference)): + rmtree(os.path.join(opt.avi_dir, opt.reference)) + + if os.path.exists(os.path.join(opt.frames_dir, opt.reference)): + rmtree(os.path.join(opt.frames_dir, opt.reference)) + + if os.path.exists(os.path.join(opt.tmp_dir, opt.reference)): + rmtree(os.path.join(opt.tmp_dir, opt.reference)) + + # ========== MAKE NEW DIRECTORIES ========== + + os.makedirs(os.path.join(opt.work_dir, opt.reference)) + os.makedirs(os.path.join(opt.crop_dir, opt.reference)) + os.makedirs(os.path.join(opt.avi_dir, opt.reference)) + os.makedirs(os.path.join(opt.frames_dir, opt.reference)) + os.makedirs(os.path.join(opt.tmp_dir, opt.reference)) + + return opt + + +def executor(opt, det_model): + # ========== ========== ========== ========== + # # EXECUTE DEMO + # ========== ========== ========== ========== + + # ========== CONVERT VIDEO AND EXTRACT FRAMES ========== + + command = "ffmpeg -y -i %s -qscale:v 2 -async 1 -r 25 %s" % ( + opt.videofile, + os.path.join(opt.avi_dir, opt.reference, "video.avi"), + ) + # output = subprocess.call(command, shell=True, stdout=None) + silent_call(command) + + command = "ffmpeg -y -i %s -qscale:v 2 -threads 1 -f image2 %s" % ( + os.path.join(opt.avi_dir, opt.reference, "video.avi"), + os.path.join(opt.frames_dir, opt.reference, "%06d.jpg"), + ) + # output = subprocess.call(command, shell=True, stdout=None) + silent_call(command) + + command = "ffmpeg -y -i %s -ac 1 -vn -acodec pcm_s16le -ar 16000 %s" % ( + os.path.join(opt.avi_dir, opt.reference, "video.avi"), + os.path.join(opt.avi_dir, opt.reference, "audio.wav"), + ) + # output = subprocess.call(command, shell=True, stdout=None) + silent_call(command) + + # ========== FACE DETECTION ========== + + faces = inference_video(opt, det_model) + + # ========== SCENE DETECTION ========== + + scene = scene_detect(opt) + + # ========== FACE TRACKING ========== + + alltracks = [] + vidtracks = [] + + for shot in scene: + if shot[1].frame_num - shot[0].frame_num >= opt.min_track: + alltracks.extend(track_shot(opt, faces[shot[0].frame_num : shot[1].frame_num])) + + # ========== FACE TRACK CROP ========== + + for ii, track in enumerate(alltracks): + vidtracks.append(crop_video(opt, track, os.path.join(opt.crop_dir, opt.reference, "%05d" % ii))) + + # ========== SAVE RESULTS ========== + + savepath = os.path.join(opt.work_dir, opt.reference, "tracks.pckl") + + with open(savepath, "wb") as fil: + pickle.dump(vidtracks, fil) + + rmtree(os.path.join(opt.tmp_dir, opt.reference)) + + +def main(data_dir, chunk, gpu_id, index): + det_model = S3FD(device=gpu_id) + + pos = index + pbar = tqdm(total=len(chunk), desc=f"Worker {index} ({gpu_id})", position=pos) + + for videofile in chunk: + reference = Path(videofile).stem + + opt = parse_args(data_dir, videofile, reference) + executor(opt, det_model) + + del opt + + pbar.update(1) + + pbar.close() + + +if __name__ == "__main__": + argparse.ArgumentParser(description="Face Detection and Tracking") + parser = argparse.ArgumentParser() + parser.add_argument("--data_dir", type=str, default="data/MAFW_GT") + parser.add_argument("--video_dir", type=str, default="/ckptstorage/zhengjunjie/data/v2sdata/clips") + parser.add_argument("--gpu_ids", type=str, default="0") + parser.add_argument("--num_workers", type=int, default=3) + + args = parser.parse_args() + + gpu_ids = [int(x) for x in args.gpu_ids.split(",")] + num_workers = args.num_workers + video_dir = args.video_dir + output_base_dir = args.data_dir + + # Ensure output directory exists + os.makedirs(output_base_dir, exist_ok=True) + + video_files = list(Path(video_dir).rglob("*.mp4")) + + chunks = np.array_split(video_files, len(gpu_ids) * num_workers) + + mp.set_start_method("spawn", force=True) + processes = [] + for idx, chunk in enumerate(chunks): + device = gpu_ids[idx % len(gpu_ids)] + + device = f"cuda:{device}" + p = mp.Process(target=main, args=(output_base_dir, chunk, device, idx)) + processes.append(p) + p.start() + + for process in processes: + process.join() + + print("All processes finished.") diff --git a/syncnet_runner.py b/syncnet_runner.py new file mode 100755 index 0000000..0597ec7 --- /dev/null +++ b/syncnet_runner.py @@ -0,0 +1,163 @@ +#!/usr/bin/python + +# -*- coding: utf-8 -*- +import warnings + +warnings.filterwarnings("ignore") + +import argparse +import glob +import os +import pickle +from pathlib import Path + +import numpy +import torch.multiprocessing as mp +from tqdm import tqdm + +from SyncNetInstance import SyncNetInstance + + +def parse_args(data_dir, videofile, reference): + opt = argparse.Namespace() + setattr(opt, "data_dir", data_dir) + setattr(opt, "videofile", videofile) + setattr(opt, "reference", reference) + + setattr(opt, "batch_size", 20) + setattr(opt, "vshift", 15) + + setattr(opt, "avi_dir", os.path.join(opt.data_dir, "pyavi")) + setattr(opt, "tmp_dir", os.path.join(opt.data_dir, "pytmp")) + setattr(opt, "work_dir", os.path.join(opt.data_dir, "pywork")) + setattr(opt, "crop_dir", os.path.join(opt.data_dir, "pycrop")) + + return opt + + +def executor(opt, s): + # ==================== LOAD MODEL AND FILE LIST ==================== + + # print("Model %s loaded." % opt.initial_model) + flist = glob.glob(os.path.join(opt.crop_dir, opt.reference, "0*.avi")) + flist.sort() + + # ==================== GET OFFSETS ==================== + + dists = [] + minvals = [] + confs = [] + + for idx, fname in enumerate(flist): + try: + offset, conf, dist, minval = s.evaluate(opt, videofile=fname) + + except Exception as e: + print(f"[ERROR] Failed to process {fname}: {e}") + continue + + dists.append(dist) + minvals.append(minval) + confs.append(conf) + + # ==================== PRINT RESULTS TO FILE ==================== + + with open(os.path.join(opt.work_dir, opt.reference, "activesd.pckl"), "wb") as fil: + pickle.dump(dists, fil) + + with open(os.path.join(opt.work_dir, opt.reference, "res.txt"), "w") as f: + f.write(f"LSE-D\t{numpy.mean(minvals, axis=0).astype(str).tolist()}") + f.write("\n") + f.write(f"LSE-C\t{numpy.mean(confs, axis=0).astype(str).tolist()}") + f.write("\n") + + +def parse_results(data_dir): + lse_ds = [] + lse_cs = [] + + for res_text_path in Path(data_dir).rglob("res.txt"): + lines = res_text_path.read_text().splitlines() + + if len(lines) < 2: + print(f"[WARNING] Invalid content in: {res_text_path}") + continue + try: + lse_d = float(lines[0].split("\t")[-1].strip()) + lse_c = float(lines[1].split("\t")[-1].strip()) + + # Exclude NaN or non-finite values + if numpy.isfinite(lse_d) and numpy.isfinite(lse_c): + lse_ds.append(lse_d) + lse_cs.append(lse_c) + else: + print(f"[WARNING] Non-finite value in: {res_text_path}") + except Exception as e: + print(f"[WARNING] Failed to parse {res_text_path}: {e}") + + if lse_ds and lse_cs: + print(f"Mean LSE-D: {numpy.mean(lse_ds):.4f}") + print(f"Mean LSE-C: {numpy.mean(lse_cs):.4f}") + else: + print("[ERROR] No valid data found.") + + +def main(data_dir, chunk, gpu_id, index): + s = SyncNetInstance(device=gpu_id) + s.loadParameters("data/syncnet_v2.model") + + pos = index + pbar = tqdm(total=len(chunk), desc=f"Worker {index} ({gpu_id})", position=pos) + + for videofile in chunk: + reference = Path(videofile).stem + + opt = parse_args(data_dir, videofile, reference) + executor(opt, s) + + del opt + + pbar.update(1) + + pbar.close() + pass + + +if __name__ == "__main__": + argparse.ArgumentParser(description="Face Detection and Tracking") + parser = argparse.ArgumentParser() + parser.add_argument("--data_dir", type=str, default="data/xxx") + parser.add_argument("--video_dir", type=str, default="/path/to/video_files") + parser.add_argument("--gpu_ids", type=str, default="0") + parser.add_argument("--num_workers", type=int, default=3) + + args = parser.parse_args() + + gpu_ids = [int(x) for x in args.gpu_ids.split(",")] + num_workers = args.num_workers + video_dir = args.video_dir + output_base_dir = args.data_dir + + # Ensure output directory exists + os.makedirs(output_base_dir, exist_ok=True) + + video_files = list(Path(video_dir).rglob("*.mp4"))[:1000] + + chunks = numpy.array_split(video_files, len(gpu_ids) * num_workers) + + mp.set_start_method("spawn", force=True) + processes = [] + for idx, chunk in enumerate(chunks): + device = gpu_ids[idx % len(gpu_ids)] + + device = f"cuda:{device}" + p = mp.Process(target=main, args=(output_base_dir, chunk, device, idx)) + processes.append(p) + p.start() + + for process in processes: + process.join() + + print("All syncnet processes finished.") + print("Parsing results...") + parse_results(output_base_dir) From c6c6110c5a85c944aa98cfba2c5021fe9456390a Mon Sep 17 00:00:00 2001 From: zodymm Date: Fri, 9 May 2025 11:37:24 +0800 Subject: [PATCH 2/3] update --- detectors/s3fd/__init__.py | 25 ++++++++++---------- detectors/s3fd/box_utils.py | 46 ++++++++++++++++++++----------------- 2 files changed, 37 insertions(+), 34 deletions(-) diff --git a/detectors/s3fd/__init__.py b/detectors/s3fd/__init__.py index d7f35e0..08db28d 100644 --- a/detectors/s3fd/__init__.py +++ b/detectors/s3fd/__init__.py @@ -1,31 +1,30 @@ import time -import numpy as np + import cv2 +import numpy as np import torch from torchvision import transforms -from .nets import S3FDNet -from .box_utils import nms_ - -PATH_WEIGHT = './detectors/s3fd/weights/sfd_face.pth' -img_mean = np.array([104., 117., 123.])[:, np.newaxis, np.newaxis].astype('float32') +from .box_utils import nms_ +from .nets import S3FDNet -class S3FD(): +PATH_WEIGHT = "./detectors/s3fd/weights/sfd_face.pth" +img_mean = np.array([104.0, 117.0, 123.0])[:, np.newaxis, np.newaxis].astype("float32") - def __init__(self, device='cuda'): +class S3FD: + def __init__(self, device="cuda"): tstamp = time.time() self.device = device - print('[S3FD] loading with', self.device) + print("[S3FD] loading with", self.device) self.net = S3FDNet(device=self.device).to(self.device) state_dict = torch.load(PATH_WEIGHT, map_location=self.device) self.net.load_state_dict(state_dict) self.net.eval() - print('[S3FD] finished loading (%.4f sec)' % (time.time() - tstamp)) - - def detect_faces(self, image, conf_th=0.8, scales=[1]): + print("[S3FD] finished loading (%.4f sec)" % (time.time() - tstamp)) + def detect_faces(self, image, conf_th=0.8, scales=[1]): w, h = image.shape[1], image.shape[0] bboxes = np.empty(shape=(0, 5)) @@ -37,7 +36,7 @@ def detect_faces(self, image, conf_th=0.8, scales=[1]): scaled_img = np.swapaxes(scaled_img, 1, 2) scaled_img = np.swapaxes(scaled_img, 1, 0) scaled_img = scaled_img[[2, 1, 0], :, :] - scaled_img = scaled_img.astype('float32') + scaled_img = scaled_img.astype("float32") scaled_img -= img_mean scaled_img = scaled_img[[2, 1, 0], :, :] x = torch.from_numpy(scaled_img).unsqueeze(0).to(self.device) diff --git a/detectors/s3fd/box_utils.py b/detectors/s3fd/box_utils.py index 0779bcd..7777fff 100644 --- a/detectors/s3fd/box_utils.py +++ b/detectors/s3fd/box_utils.py @@ -1,5 +1,6 @@ -import numpy as np from itertools import product as product + +import numpy as np import torch from torch.autograd import Function @@ -35,7 +36,7 @@ def nms_(dets, thresh): inds = np.where(ovr <= thresh)[0] order = order[inds + 1] - return np.array(keep).astype(np.int) + return np.array(keep).astype(int) def decode(loc, priors, variances): @@ -51,9 +52,13 @@ def decode(loc, priors, variances): decoded bounding box predictions """ - boxes = torch.cat(( - priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:], - priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1) + boxes = torch.cat( + ( + priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:], + priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1]), + ), + 1, + ) boxes[:, :2] -= boxes[:, 2:] / 2 boxes[:, 2:] += boxes[:, :2] return boxes @@ -127,11 +132,9 @@ def nms(boxes, scores, overlap=0.5, top_k=200): class Detect(object): - - def __init__(self, num_classes=2, - top_k=750, nms_thresh=0.3, conf_thresh=0.05, - variance=[0.1, 0.2], nms_top_k=5000): - + def __init__( + self, num_classes=2, top_k=750, nms_thresh=0.3, conf_thresh=0.05, variance=[0.1, 0.2], nms_top_k=5000 + ): self.num_classes = num_classes self.top_k = top_k self.nms_thresh = nms_thresh @@ -140,7 +143,6 @@ def __init__(self, num_classes=2, self.nms_top_k = nms_top_k def forward(self, loc_data, conf_data, prior_data): - num = loc_data.size(0) num_priors = prior_data.size(0) @@ -160,7 +162,7 @@ def forward(self, loc_data, conf_data, prior_data): for cl in range(1, self.num_classes): c_mask = conf_scores[cl].gt(self.conf_thresh) scores = conf_scores[cl][c_mask] - + if scores.dim() == 0: continue l_mask = c_mask.unsqueeze(1).expand_as(boxes) @@ -174,13 +176,15 @@ def forward(self, loc_data, conf_data, prior_data): class PriorBox(object): - - def __init__(self, input_size, feature_maps, - variance=[0.1, 0.2], - min_sizes=[16, 32, 64, 128, 256, 512], - steps=[4, 8, 16, 32, 64, 128], - clip=False): - + def __init__( + self, + input_size, + feature_maps, + variance=[0.1, 0.2], + min_sizes=[16, 32, 64, 128, 256, 512], + steps=[4, 8, 16, 32, 64, 128], + clip=False, + ): super(PriorBox, self).__init__() self.imh = input_size[0] @@ -210,8 +214,8 @@ def forward(self): mean += [cx, cy, s_kw, s_kh] output = torch.FloatTensor(mean).view(-1, 4) - + if self.clip: output.clamp_(max=1, min=0) - + return output From 3fe9d520f4f347bd08628261d77cac790eb13019 Mon Sep 17 00:00:00 2001 From: zodymm Date: Fri, 9 May 2025 11:37:53 +0800 Subject: [PATCH 3/3] add run.sh --- run.sh | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 run.sh diff --git a/run.sh b/run.sh new file mode 100644 index 0000000..42423a2 --- /dev/null +++ b/run.sh @@ -0,0 +1,2 @@ +python crop_runner.py --output_base_dir data/xxx --video_dir /path/to/video_dir +python syncnet_runner.py --data_dir data/xxx --video_dir /path/to/video_dir