From 065774a5016d920d3ecdc3c9e9423f6d13865cf7 Mon Sep 17 00:00:00 2001
From: Zangwei Zheng <zangwei@comp.nus.edu.sg>
Date: Tue, 2 Apr 2024 16:55:03 +0800
Subject: [PATCH] [feat] resize images

---
 tools/datasets/README.md             | 18 ++++++
 tools/datasets/resize.py             | 87 ++++++++++++++++++++++++++++
 tools/scoring/aesthetic/inference.py |  2 +-
 3 files changed, 106 insertions(+), 1 deletion(-)
 create mode 100644 tools/datasets/resize.py

diff --git a/tools/datasets/README.md b/tools/datasets/README.md
index 0d4d721..6b34287 100644
--- a/tools/datasets/README.md
+++ b/tools/datasets/README.md
@@ -1,5 +1,15 @@
 # Dataset Management
 
+- [Dataset Management](#dataset-management)
+  - [Dataset Format](#dataset-format)
+  - [Dataset to CSV](#dataset-to-csv)
+  - [Manage datasets](#manage-datasets)
+    - [Requirement](#requirement)
+    - [Usage](#usage)
+  - [Analyze datasets](#analyze-datasets)
+  - [Frame extraction speed](#frame-extraction-speed)
+
+
 After preparing the raw dataset according to the [instructions](/docs/datasets.md), you can use the following commands to manage the dataset.
 
 ## Dataset Format
@@ -175,6 +185,14 @@ data.hist(column="resolution")
 plt.savefig('info.jpg')
 ```
 
+## Resize datasets
+
+Sometimes you may need to resize the images or videos to a specific resolution. You can use the following commands to resize the dataset:
+
+```bash
+python -m tools.datasets.resize meta.csv /path/to/raw/data /path/to/new/data --length 2160
+```
+
 ## Frame extraction speed
 
 We use three libraries to extract frames from videos: `opencv`, `pyav` and `decord`. Our benchmark results of loading 256 video's middle frames are as follows:
diff --git a/tools/datasets/resize.py b/tools/datasets/resize.py
new file mode 100644
index 0000000..352be53
--- /dev/null
+++ b/tools/datasets/resize.py
@@ -0,0 +1,87 @@
+import argparse
+import os
+
+import cv2
+import pandas as pd
+from tqdm import tqdm
+
+tqdm.pandas()
+
+try:
+    from pandarallel import pandarallel
+
+    pandarallel.initialize(progress_bar=True)
+    pandas_has_parallel = True
+except ImportError:
+    pandas_has_parallel = False
+
+
+def apply(df, func):
+    if pandas_has_parallel:
+        return df.parallel_apply(func)
+    return df.progress_apply(func)
+
+
+IMG_EXTENSIONS = (
+    ".jpg",
+    ".jpeg",
+    ".png",
+    ".ppm",
+    ".bmp",
+    ".pgm",
+    ".tif",
+    ".tiff",
+    ".webp",
+)
+
+
+def get_new_path(path, input_dir, output):
+    path_new = os.path.join(output, os.path.relpath(path, input_dir))
+    os.makedirs(os.path.dirname(path_new), exist_ok=True)
+    return path_new
+
+
+def resize(path, length, input_dir, output):
+    path_new = get_new_path(path, input_dir, output)
+    ext = os.path.splitext(path)[1].lower()
+    if ext in IMG_EXTENSIONS:
+        img = cv2.imread(path)
+        h, w = img.shape[:2]
+        if min(h, w) > length:
+            if h > w:
+                new_h = length
+                new_w = int(w * new_h / h)
+            else:
+                new_w = length
+                new_h = int(h * new_w / w)
+            img = cv2.resize(img, (new_w, new_h))
+        cv2.imwrite(path_new, img)
+    else:
+        pass
+    return path_new
+
+
+def main(args):
+    data = pd.read_csv(args.input)
+    data["path"] = apply(data["path"], lambda x: resize(x, args.length, args.input_dir, args.output))
+    output_csv = args.input.replace(".csv", f"_resized{args.length}.csv")
+    data.to_csv(output_csv, index=False)
+    print(f"Saved to {output_csv}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("input", type=str)
+    parser.add_argument("input_dir", type=str)
+    parser.add_argument("output", type=str)
+    parser.add_argument("--disable-parallel", action="store_true")
+    parser.add_argument("--length", type=int, default=2160)
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    if args.disable_parallel:
+        pandas_has_parallel = False
+    main(args)
diff --git a/tools/scoring/aesthetic/inference.py b/tools/scoring/aesthetic/inference.py
index f28c4bc..d73a133 100644
--- a/tools/scoring/aesthetic/inference.py
+++ b/tools/scoring/aesthetic/inference.py
@@ -130,7 +130,7 @@ def main(args):
         dataset.data.loc[index : index + len(scores_np) - 1, "aes"] = scores_np
         index += len(images)
     dataset.data.to_csv(output_file, index=False)
-    print(f"New meta with aesthetic scores saved to \'{output_file}\'.")
+    print(f"New meta with aesthetic scores saved to '{output_file}'.")
 
 
 if __name__ == "__main__":