Change drive process for saving images, change dataset root

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,23 +76,19 @@ 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)
 
             embedder.get_pose_embedding(data_dict)
             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

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.: