From 3296b7cb5d58fb9bae6b1af587074964205cd65f Mon Sep 17 00:00:00 2001
From: Denis Leonov <denleon@phygitalism.com>
Date: Thu, 4 Mar 2021 15:00:05 +0300
Subject: [PATCH] Change drive process for saving images, change dataset root

---
 drive.py                    | 20 ++++++++------------
 utils/preprocess_dataset.sh |  4 ++--
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/drive.py b/drive.py
index 61a2ffa..aea2876 100644
--- a/drive.py
+++ b/drive.py
@@ -1,9 +1,9 @@
 import numpy as np
 import cv2
 import torch
+import os 
 torch.set_grad_enabled(False)
 
-from utils.crop_as_in_dataset import ImageWriter
 from utils import utils
 
 from pathlib import Path
@@ -76,11 +76,8 @@ def string_to_valid_filename(x):
         from dataloaders.dataloader import Dataloader
         logger.info(f"Loading dataloader '{saved_args.dataloader}'")
         dataloader = Dataloader(saved_args.dataloader).get_dataloader(saved_args, part='val', phase='val')
-
-        current_output_path = (args.destination / string_to_valid_filename(driver_images_path)).with_suffix('.mp4')
-        current_output_path.parent.mkdir(parents=True, exist_ok=True)
-        image_writer = ImageWriter.get_image_writer(current_output_path)
-
+        
+        _ind = 0
         for data_dict, _ in tqdm(dataloader):
             utils.dict_to_device(data_dict, device)
 
@@ -88,11 +85,10 @@ def string_to_valid_filename(x):
             generator(data_dict)
 
             def torch_to_opencv(image):
-                image = image.permute(1,2,0).clamp_(0, 1).mul_(255).cpu().byte().numpy()
-                return cv2.cvtColor(image, cv2.COLOR_RGB2BGR, dst=image)
+                image = image.permute(1, 2, 0).clamp_(0, 1).mul_(255).cpu().byte().numpy()
+                return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
             result = torch_to_opencv(data_dict['fake_rgbs'][0])
-            pose_driver = torch_to_opencv(data_dict['pose_input_rgbs'][0, 0])
-
-            frame_grid = np.concatenate((pose_driver, result), axis=1)
-            image_writer.add(frame_grid)
+            filename = os.path.join(args.destination, str(_ind).zfill(3) + '.png')
+            cv2.imwrite(filename, result)
+            _ind += 1
\ No newline at end of file
diff --git a/utils/preprocess_dataset.sh b/utils/preprocess_dataset.sh
index 3095585..25f88e0 100644
--- a/utils/preprocess_dataset.sh
+++ b/utils/preprocess_dataset.sh
@@ -33,14 +33,14 @@
 set -e
 
 # Please use an ABSOLUTE path here!
-DATASET_ROOT="/Vol1/dbstore/datasets/violet/VoxCeleb2_test_finetuning"
+DATASET_ROOT="/home/app/latent-pose-reenactment/dataset"
 
 # echo "Unnamed: 0,path" > $dataset_dir/split.csv
 
 # Initialize `IDENTITIES` -- the list of folders (paths relative to $DATASET_ROOT/images
 # or $DATASET_ROOT/videos), each containing raw images or one video of some person.
 cd "$DATASET_ROOT/images" # or e.g. `"$DATASET_ROOT/videos"`
-IDENTITIES=(*) # or `(*/*)`, or whatever else
+IDENTITIES=('someone') # or `(*/*)`, or whatever else
 cd -
 
 # Alternatively, you can specify them manually, e.g.: