order by saliency

Author: Delphboy

.vscode/launch.json

@@ -4,11 +4,25 @@
     // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
     "version": "0.2.0",
     "configurations": [
+        {
+            "name": "Sort by Saliency",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "feature_order/saliency.py",
+            "args": [
+                "--karpathy_json_file", "/homes/hps01/dataset_coco_10.json",
+                "--image_dir", "/import/gameai-01/eey362/datasets/coco/images/",
+                "--feature_dir", "/import/gameai-01/eey362/VSUA-Captioning/data/cocobu_att/",
+                "--bbox_dir", "/import/gameai-01/eey362/VSUA-Captioning/data/cocobu_box/",
+                "--output_dir", "/import/gameai-01/eey362/datasets/coco/butd_ordered_saliency",
+            ],
+            "console": "integratedTerminal"
+        },
         {
             "name": "Sort BUTD by bbox",
             "type": "debugpy",
             "request": "launch",
-            "program": "feature_order/feature_order.py",
+            "program": "feature_order/bbox_size.py",
             "args": [
                 "--bbox_dir", "/import/gameai-01/eey362/VSUA-Captioning/data/cocobu_box/",
                 "--feature_dir", "/import/gameai-01/eey362/VSUA-Captioning/data/cocobu_att/",

feature_order/feature_order.py renamed to feature_order/bbox_size.py

@@ -1,14 +1,8 @@
 import os
 import argparse
 import numpy as np
-from tqdm import tqdm
-from multiprocessing import Pool
 
-def load_feature(feature_dir):
-    return np.load(feature_dir)["feat"]
-
-def load_bbox(bbox_dir):
-    return np.load(bbox_dir) # [B, 4]
+from helpers import load_feature, load_bbox, save_feature, process_parallel
 
 def caculate_area(bbox):
     x1 = bbox[:, 0]
@@ -24,11 +18,6 @@ def sort_feature_by_area(features, bbox):
     features = features[sort_index]
     return features
 
-def save_feature(output_dir, name, features):
-    path = os.path.join(output_dir, f"{name}.npz")
-    np.savez_compressed(path, feat=features)
-
-
 def process_individual(feature_dir, bbox_dir, output_dir, file):
     feature_path = os.path.join(feature_dir, file)
     features = load_feature(feature_path)
@@ -41,16 +30,6 @@ def process_individual(feature_dir, bbox_dir, output_dir, file):
     save_feature(output_dir, file.split(".")[0], sorted_features)
 
 
-# Using multiprocess, rewrite the process function so that it takes advantage of parallelism
-# and processes the features in parallel without tqdm progress bar. Instead, use Counter to keep track of the progress.
-# and print the progress every 1000 images.
-def process_parallel(feature_dir, bbox_dir, output_dir):
-    files = os.listdir(feature_dir)
-    with Pool() as p:
-        p.starmap(process_individual, [(feature_dir, bbox_dir, output_dir, file) for file in files])
-
-
-
 if __name__ == "__main__":
     args = argparse.ArgumentParser()
     args.add_argument("--feature_dir", type=str, required=True)
@@ -61,4 +40,4 @@ def process_parallel(feature_dir, bbox_dir, output_dir):
 
     os.makedirs(args.output_dir, exist_ok=True)
 
-    process_parallel(args.feature_dir, args.bbox_dir, args.output_dir)
+    process_parallel(args.feature_dir, args.bbox_dir, args.output_dir, process_individual)

feature_order/coco.py

@@ -0,0 +1,66 @@
+import json
+import os
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+from PIL import Image
+
+class CocoButdFeatures(Dataset):
+    def __init__(
+        self,
+        captions_file: str,
+        image_dir: str,
+        feature_dir: str,
+        bbox_dir: str,
+    ):
+        self.captions_file = captions_file
+        self.image_dir = image_dir
+        self.feature_dir = feature_dir
+        self.bbox_dir = bbox_dir
+
+        with open(self.captions_file, "r") as f:
+            self.captions_file_data = json.load(f)
+
+        self.image_locations = []
+        self.feature_locations = []
+        self.box_locations = []
+
+        for image_data in self.captions_file_data["images"]:
+            self.image_locations.append(
+                os.path.join(
+                    self.image_dir,
+                    image_data["filepath"],
+                    image_data["filename"],
+                )
+            )
+
+            feat_path = os.path.join(
+                    self.feature_dir,
+                    (f"{image_data['cocoid']}.npz"),
+                )
+            self.feature_locations.append(feat_path)
+
+            bbox_path = os.path.join(
+                self.bbox_dir,
+                (f"{image_data['cocoid']}.npy"),
+            )
+            self.box_locations.append(bbox_path)
+
+    def __getitem__(self, index):
+        image = Image.open(self.image_locations[index])
+        image = np.array(image)
+
+        features = np.load(self.feature_locations[index])["feat"]
+        
+        box_loc = self.box_locations[index]
+        boxes = np.load(box_loc)
+        boxes = torch.from_numpy(boxes).type(torch.float32)
+
+        name = box_loc.split("/")[-1].split(".")[0]
+
+        return name, image, features, boxes
+
+    def __len__(self):
+        return len(self.box_locations)

feature_order/helpers.py

@@ -0,0 +1,24 @@
+import numpy as np
+import os
+from multiprocessing import Pool
+
+def load_feature(feature_dir):
+    return np.load(feature_dir)["feat"]
+
+
+def load_bbox(bbox_dir):
+    return np.load(bbox_dir)  # [B, 4]
+
+
+def save_feature(output_dir, name, features):
+    path = os.path.join(output_dir, f"{name}.npz")
+    np.savez_compressed(path, feat=features)
+
+
+def process_parallel(feature_dir, bbox_dir, output_dir, function):
+    files = os.listdir(feature_dir)
+    with Pool() as p:
+        p.starmap(
+            function,
+            [(feature_dir, bbox_dir, output_dir, file) for file in files],
+        )

feature_order/saliency.py

@@ -0,0 +1,55 @@
+import argparse
+from multiprocessing import Pool
+import os
+
+import cv2
+import numpy as np
+
+from coco import CocoButdFeatures
+from helpers import save_feature
+
+
+def get_order_index_from_saliency(image, bboxes):
+    saliency = cv2.saliency.StaticSaliencySpectralResidual_create()
+    (success, saliencyMap) = saliency.computeSaliency(image)
+    saliency_scores = []
+    for bbox in bboxes:
+        x1, y1, x2, y2 = bbox
+        saliency_score = np.mean(saliencyMap[int(y1):int(y2), int(x1):int(x2)])
+        saliency_scores.append(saliency_score)
+
+    return np.argsort(saliency_scores)[::-1]
+
+
+def process_individual(dataset, index, output_dir):
+    name, image, features, boxes = dataset.__getitem__(index)
+    order = get_order_index_from_saliency(image, boxes)
+    ordered_features = features[order]
+    save_feature(output_dir, name, ordered_features)
+
+
+def process_parallel(output_dir, dataset, function):
+    with Pool() as p:
+        p.starmap(function, [(dataset, i, output_dir) for i in range(len(dataset))])
+
+
+if __name__ == "__main__":
+    args = argparse.ArgumentParser()
+    args.add_argument("--karpathy_json_file", type=str, required=True)
+    args.add_argument("--image_dir", type=str, required=True)
+    args.add_argument("--feature_dir", type=str, required=True)
+    args.add_argument("--output_dir", type=str, required=True)
+    args.add_argument("--bbox_dir", type=str, required=True)
+
+    args = args.parse_args()
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    dataset = CocoButdFeatures(
+        captions_file=args.karpathy_json_file,
+        image_dir=args.image_dir,
+        feature_dir=args.feature_dir,
+        bbox_dir=args.bbox_dir,
+    )
+
+    process_parallel(args.output_dir, dataset, process_individual)