fix eval script

evaluate.py

@@ -56,10 +56,9 @@ def compute_fid(img1, img2):
     return fid
 
 
-with open('metrics.json', 'r') as file:
-    metrics = json.load(file)
 
-def get_score(item, image_paths, video_path, train_steps=100, inference_steps=10, fps=12, bg_remove=False):
+
+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:
@@ -67,6 +66,9 @@ def get_score(item, image_paths, video_path, train_steps=100, inference_steps=10
         images.append([img])
     
     gt_frames = extract_frames(video_path, fps)
+    #gt_frames = gt_frames[:2]
+    for f in gt_frames:
+        f.thumbnail((512,512))
     
     os.makedirs('out/'+item, exist_ok=True)
     
@@ -76,8 +78,6 @@ def get_score(item, image_paths, video_path, train_steps=100, inference_steps=10
     
     #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)
-    
-    print(results)
 
     for i, result in enumerate(results):
         result.save("out/"+item+"/result_"+str(i)+".png")
@@ -134,6 +134,8 @@ def get_score(item, image_paths, video_path, train_steps=100, inference_steps=10
     metrics[item]['base']['lpips'] = {'avg': sum(lpips2)/len(lpips2), 'vals': lpips2}
     metrics[item]['base']['fid'] = {'avg': sum(fid2)/len(fid2), 'vals': fid2}
     
+    #print(metrics)
+
     with open('metrics.json', "w", encoding="utf-8") as json_file:
         json.dump(metrics, json_file, ensure_ascii=False, indent=4)
     
@@ -154,30 +156,33 @@ def get_files(directory_path):
 def run_evaluate():
     snapshot_download(repo_id="acmyu/KeyframesAI-eval", local_dir="test", repo_type="dataset")
 
+    with open('metrics.json', 'r') as file:
+        metrics = json.load(file)
+
     items = os.listdir('test') 
     items = [it for it in items if not it[0]=='.' and not os.path.isfile('test/'+it)]
     print(items)
-    items = ['sidewalk'] #['sidewalk', 'aaa', 'azri', 'dead', 'frankgirl', 'kobold', 'ramona', 'renee', 'walk', 'woody']
+    #items = ['sidewalk'] #['sidewalk', 'aaa', 'azri', 'dead', 'frankgirl', 'kobold', 'ramona', 'renee', 'walk', 'woody']
     
     for item in items: 
         if item in metrics:
             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])
-                #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:
-            print("Error", 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:
+        #    print("Error", item)
 
 
     ssim = []

main.py

@@ -1271,14 +1271,17 @@ def run_eval(images_orig, video_path, train_steps=100, inference_steps=10, fps=1
     in_img, in_pose, train_imgs, train_poses = prepare_inputs_train(images, bg_remove, dwpose, rembg_session)
     in_img, target_poses, in_pose, _, _ = prepare_inputs_inference(in_img, video_path, [], fps, dwpose, rembg_session, bg_remove, resize_inputs, is_app)
     
-    finetune = True
-    train(modelId, in_img, in_pose, train_imgs, train_poses, train_steps, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet, finetune, is_app)
-    results = inference(modelId, in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
-
+    #target_poses = target_poses[:2]
+    #train_steps = 3
+    
     finetune = False
     train(modelId, in_img, in_pose, train_imgs, train_poses, train_steps, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet, finetune, is_app)
     results_base = inference(modelId, in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
 
+    finetune = True
+    train(modelId, in_img, in_pose, train_imgs, train_poses, train_steps, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet, finetune, is_app)
+    results = inference(modelId, in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
+
     gc.collect()
     torch.cuda.empty_cache()