get fvd metric

evaluate.py

@@ -8,14 +8,26 @@ import torch
 import torchvision.transforms as transforms
 import lpips
 from pytorch_fid.fid_score import calculate_fid_given_paths
+from cdfvd import fvd
 import os
 import json
+import cv2
 from huggingface_hub import snapshot_download
 
 # Convert PIL to numpy
 def pil_to_np(img):
     return np.array(img).astype(np.float32) / 255.0
 
+def save_mp4(images, name):
+    width, height = images[0].size
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v') # Codec for MP4
+    video = cv2.VideoWriter(name, fourcc, 12, (width, height))
+
+    for image in images:
+        img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+        video.write(img)
+    video.release()
+
 # SSIM
 def compute_ssim(img1, img2):
     img1_np = pil_to_np(img1)
@@ -46,53 +58,96 @@ def compute_lpips(img1, img2):
     img2_tensor = lpips_transform(img2).unsqueeze(0)
     return lpips_model(img1_tensor, img2_tensor).item()
 
+
+def trans(x):
+    # if greyscale images add channel
+    if x.shape[-3] == 1:
+        x = x.repeat(1, 1, 3, 1, 1)
+
+    # permute BTCHW -> BCTHW
+    x = x.permute(0, 2, 1, 3, 4) 
+
+    return x
+
+
+def compute_fvd(item, gt_imgs, results):
+    os.makedirs('temp/gt', exist_ok=True)
+    os.makedirs('temp/result', exist_ok=True)
+
+    #save_mp4(gt_imgs, "temp/gt/gt.mp4")
+    #save_mp4(results, "temp/result/result.mp4")
+
+    evaluator = fvd.cdfvd('i3d', ckpt_path=None, device='cuda', n_real=1, n_fake=1)
+    evaluator.compute_real_stats(evaluator.load_videos('temp/gt', data_type='video_folder'))
+    evaluator.compute_fake_stats(evaluator.load_videos('temp/result', data_type='video_folder'))
+    score = evaluator.compute_fvd_from_stats()
+    evaluator.offload_model_to_cpu()
+    print(score)
+    return score
+
+
+def compute_fidx(item, gt_imgs, results):
+    os.makedirs('temp/'+item+'_gt', exist_ok=True)
+    os.makedirs('temp/'+item, exist_ok=True)
+    c = 0
+    for img in gt_imgs:
+        img.save('temp/'+item+'_gt/'+str(c)+'.png')
+        c = c+1
+    c = 0
+    for img in gt_imgs:
+        img.save('temp/'+item+'/'+str(c)+'.png')
+        c = c+1
+
+    fid = calculate_fid_given_paths(['temp/'+item+'_gt', 'temp/'+item], batch_size=8, device='cuda', dims=2048)
+    return fid
+
 # FID: Save images to temp folders for FID calculation
 def compute_fid(img1, img2):
     os.makedirs('temp/img1', exist_ok=True)
     os.makedirs('temp/img2', exist_ok=True)
     img1.save('temp/img1/0.png')
     img2.save('temp/img2/0.png')
-    fid = calculate_fid_given_paths(['temp/img1', 'temp/img2'], batch_size=1, device='cpu', dims=2048)
+    fid = calculate_fid_given_paths(['temp/img1', 'temp/img2'], batch_size=1, device='cuda', dims=2048)
     return fid
 
 
-
-
 def get_score(item, image_paths, video_path, metrics, train_steps=100, inference_steps=10, fps=12, bg_remove=False):
-
     images = []
     for path in image_paths:
         img = Image.open(path)
         images.append([img])
-    
-    gt_frames = extract_frames(video_path, fps)
-    gt_frames = gt_frames[:1000]
-    for f in gt_frames:
-        f.thumbnail((512,512))
-    
-    os.makedirs('out/'+item, exist_ok=True)
-    
-    
-    for i, frame in enumerate(gt_frames):
-        frame.save("out/"+item+"/frame_"+str(i)+".png")
-    
-    #results = run(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_remove=False, finetune=True)
-    results, results_base = run_eval(images, video_path, train_steps=100, inference_steps=10, fps=12, modelId="fine_tuned_pcdms", img_width=1920, img_height=1080, bg_remove=False, resize_inputs=False)
 
-    for i, result in enumerate(results):
-        result.save("out/"+item+"/result_"+str(i)+".png")
-        
-    for i, result in enumerate(results_base):
-        result.save("out/"+item+"/base_"+str(i)+".png")
-    
-    """
-    img1=gt_frames[0]
-    img2=Image.open("out/base_0.png")
-    print("SSIM:", compute_ssim(img1, img2))
-    print("PSNR:", compute_psnr(img1, img2))
-    print("LPIPS:", compute_lpips(img1, img2))
-    print("FID:", compute_fid(img1, img2))
-    """
+    results = []
+    results_base = []
+    gt_frames = []
+    if os.path.isdir('out/'+item):
+        for filename in os.listdir('out/'+item):
+            img = Image.open('out/'+item+'/'+filename)
+            if filename.startswith('result_'):
+                results.append(img)
+            elif filename.startswith('base_'):
+                results_base.append(img)
+            elif filename.startswith('frame_'):
+                gt_frames.append(img)
+    else:
+        gt_frames = extract_frames(video_path, fps)
+        gt_frames = gt_frames[:200]
+        for f in gt_frames:
+            f.thumbnail((512,512))
+
+        #results = run(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_remove=False, finetune=True)
+        results, results_base = run_eval(images, video_path, train_steps=100, inference_steps=10, fps=12, modelId="fine_tuned_pcdms", img_width=1920, img_height=1080, bg_remove=False, resize_inputs=False)
+
+        os.makedirs('out/'+item, exist_ok=True)
+
+        for i, frame in enumerate(gt_frames):
+            frame.save("out/"+item+"/frame_"+str(i)+".png")
+
+        for i, result in enumerate(results):
+            result.save("out/"+item+"/result_"+str(i)+".png")
+            
+        for i, result in enumerate(results_base):
+            result.save("out/"+item+"/base_"+str(i)+".png")
     
     ssim = []
     psnr = []
@@ -102,37 +157,50 @@ def get_score(item, image_paths, video_path, metrics, train_steps=100, inference
     psnr2 = []
     lpips2 = []
     fid2 = []
+    c = 0
+    #print(len(gt_frames), len(results), len(results_base))
     for gt, result, base in zip(gt_frames, results, results_base):
         ssim.append(float(compute_ssim(gt, result)))
         psnr.append(float(compute_psnr(gt, result)))
         lpips.append(float(compute_lpips(gt, result)))
-        fid.append(float(compute_fid(gt, result)))
         
         ssim2.append(float(compute_ssim(gt, base)))
         psnr2.append(float(compute_psnr(gt, base)))
         lpips2.append(float(compute_lpips(gt, base)))
-        fid2.append(float(compute_fid(gt, base)))
         
+        if c<50:
+            print(c)
+            fid.append(float(compute_fid(gt, result)))
+            fid2.append(float(compute_fid(gt, base)))
+        c = c+1
+        
+    fvd = float(compute_fvd(item, gt_frames, results))
+    fvd2 = float(compute_fvd(item, gt_frames, results_base))
+    
     
     print("SSIM:", sum(ssim)/len(ssim))
     print("PSNR:", sum(psnr)/len(psnr))
     print("LPIPS:", sum(lpips)/len(lpips))
     print("FID:", sum(fid)/len(fid))
+    print("FVD:", fvd)
     print('baseline:')
     print("SSIM:", sum(ssim2)/len(ssim2))
     print("PSNR:", sum(psnr2)/len(psnr2))
     print("LPIPS:", sum(lpips2)/len(lpips2))
     print("FID:", sum(fid2)/len(fid2))
+    print("FVD:", fvd2)
     
-    metrics[item] = {'ft': {}, 'base': {}}
+    metrics[item] = {'ft': {}, 'base': {}, 'frames': len(gt_frames), 'complexity': len(images)}
     metrics[item]['ft']['ssim'] = {'avg': sum(ssim)/len(ssim), 'vals': ssim}
     metrics[item]['ft']['psnr'] = {'avg': sum(psnr)/len(psnr), 'vals': psnr}
     metrics[item]['ft']['lpips'] = {'avg': sum(lpips)/len(lpips), 'vals': lpips}
     metrics[item]['ft']['fid'] = {'avg': sum(fid)/len(fid), 'vals': fid}
+    metrics[item]['ft']['fvd'] = fvd
     metrics[item]['base']['ssim'] = {'avg': sum(ssim2)/len(ssim2), 'vals': ssim2}
     metrics[item]['base']['psnr'] = {'avg': sum(psnr2)/len(psnr2), 'vals': psnr2}
     metrics[item]['base']['lpips'] = {'avg': sum(lpips2)/len(lpips2), 'vals': lpips2}
     metrics[item]['base']['fid'] = {'avg': sum(fid2)/len(fid2), 'vals': fid2}
+    metrics[item]['base']['fvd'] = fvd2
     
     #print(metrics)
 
@@ -154,6 +222,7 @@ def get_files(directory_path):
 
 
 def run_evaluate():
+    print("run_evaluate")
     snapshot_download(repo_id="acmyu/KeyframesAI-eval", local_dir="test", repo_type="dataset")
 
     with open('metrics.json', 'r') as file:
@@ -169,20 +238,20 @@ def run_evaluate():
             continue
         print(item)
         
-        try:
-            files = get_files('test/'+item)
-            images = list(filter(lambda x: not x.endswith('.mp4'), files))
-            images = ['test/'+item+'/'+img for img in images]
-            videos = [x for x in files if x.endswith('.mp4')]
-            print(images, videos)
-
-            if len(videos) == 1:
-                get_score(item, images, 'test/'+item+'/'+videos[0], metrics)
-                #get_score(item, ['test/'+item+'/1.jpg', 'test/'+item+'/2.jpg', 'test/'+item+'/3.jpg'], 'test/'+item+'/v.mp4')
-            else:
-                print('Error: mp4 not found')
-        except Exception as e:
-            print("Error", item, e)
+        #try:
+        files = get_files('test/'+item)
+        images = list(filter(lambda x: not x.endswith('.mp4'), files))
+        images = ['test/'+item+'/'+img for img in images]
+        videos = [x for x in files if x.endswith('.mp4')]
+        print(images, videos)
+
+        if len(videos) == 1:
+            get_score(item, images, 'test/'+item+'/'+videos[0], metrics)
+            #get_score(item, ['test/'+item+'/1.jpg', 'test/'+item+'/2.jpg', 'test/'+item+'/3.jpg'], 'test/'+item+'/v.mp4')
+        else:
+            print('Error: mp4 not found')
+        #except Exception as e:
+        #    print("Error", item, e)
 
 
     ssim = []

main.py

@@ -85,7 +85,7 @@ debug = False
 save_model = True
 should_gen_vid = False
 max_batch_size = 8
-max_frame_count = 1000
+max_frame_count = 200
 
 def save_temp_imgs(imgs):
     os.makedirs('temp', exist_ok=True)

requirements.txt

@@ -23,4 +23,6 @@ spaces
 matplotlib
 
 lpips
-pytorch-fid
+pytorch-fid
+cd-fvd
+av