Spaces:
Runtime error
Runtime error
import argparse | |
import os | |
os.environ["MKL_NUM_THREADS"] = "1" | |
os.environ["NUMEXPR_NUM_THREADS"] = "1" | |
os.environ["OMP_NUM_THREADS"] = "1" | |
# from skimage.measure import compare_ssim | |
from skimage.metrics import structural_similarity as ssim | |
from functools import partial | |
from multiprocessing.pool import Pool | |
from tqdm import tqdm | |
from utils import get_original_with_fakes | |
import cv2 | |
cv2.ocl.setUseOpenCL(False) | |
cv2.setNumThreads(0) | |
import numpy as np | |
cache = {} | |
def save_diffs(pair, root_dir): | |
ori_id, fake_id = pair | |
ori_dir = os.path.join(root_dir, "crops", ori_id) | |
fake_dir = os.path.join(root_dir, "crops", fake_id) | |
diff_dir = os.path.join(root_dir, "diffs", fake_id) | |
os.makedirs(diff_dir, exist_ok=True) | |
for frame in range(320): | |
if frame % 10 != 0: | |
continue | |
for actor in range(2): | |
image_id = "{}_{}.png".format(frame, actor) | |
diff_image_id = "{}_{}_diff.png".format(frame, actor) | |
ori_path = os.path.join(ori_dir, image_id) | |
fake_path = os.path.join(fake_dir, image_id) | |
diff_path = os.path.join(diff_dir, diff_image_id) | |
if os.path.exists(ori_path) and os.path.exists(fake_path): | |
img1 = cv2.imread(ori_path, cv2.IMREAD_COLOR) | |
img2 = cv2.imread(fake_path, cv2.IMREAD_COLOR) | |
try: | |
d, a = ssim(img1, img2, multichannel=True, full=True) | |
a = 1 - a | |
diff = (a * 255).astype(np.uint8) | |
diff = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY) | |
cv2.imwrite(diff_path, diff) | |
except: | |
pass | |
def parse_args(): | |
parser = argparse.ArgumentParser( | |
description="Extract image diffs") | |
parser.add_argument("--root-dir", help="root directory", default="/mnt/sota/datasets/deepfake") | |
args = parser.parse_args() | |
return args | |
def main(): | |
args = parse_args() | |
pairs = get_original_with_fakes(args.root_dir) | |
os.makedirs(os.path.join(args.root_dir, "diffs"), exist_ok=True) | |
with Pool(processes=os.cpu_count() - 2) as p: | |
with tqdm(total=len(pairs)) as pbar: | |
func = partial(save_diffs, root_dir=args.root_dir) | |
for v in p.imap_unordered(func, pairs): | |
pbar.update() | |
if __name__ == '__main__': | |
main() | |
z=2 | |