We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Minor bug in object detections in extract_features_vmb.py
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

I believe this was introduced by the #127, but there's a slight numpy indexing issue that causes the objects to still be detected incorrectly. 

The line in question is [here](https://github.com/facebookresearch/pythia/blob/12f67cd4f67499814bb0b3665ff14dd635800f63/pythia/scripts/features/extract_features_vmb.py#L165). It currently reads

```python
objects = torch.argmax(scores[keep_boxes][start_index:], dim=1)
```

However, `scores` is a tensor of `[num_objects,  object_classes]`, so `start_index` should be indexing the second dimension. The updated line should be

```python
objects = torch.argmax(scores[keep_boxes][:, start_index:], dim=1)
```

I can submit a pull request.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pythia/scripts/features/extract_features_vmb.py`
Content:
```
1 # Requires vqa-maskrcnn-benchmark to be built and installed
2 # Category mapping for visual genome can be downloaded from
3 # https://dl.fbaipublicfiles.com/pythia/data/visual_genome_categories.json
4 # When the --background flag is set, the index saved with key "objects" in
5 # info_list will be +1 of the Visual Genome category mapping above and 0
6 # is the background class. When the --background flag is not set, the
7 # index saved with key "objects" in info list will match the Visual Genome
8 # category mapping.
9 import argparse
10 import glob
11 import os
12 
13 import cv2
14 import numpy as np
15 import torch
16 from PIL import Image
17 
18 from maskrcnn_benchmark.config import cfg
19 from maskrcnn_benchmark.layers import nms
20 from maskrcnn_benchmark.modeling.detector import build_detection_model
21 from maskrcnn_benchmark.structures.image_list import to_image_list
22 from maskrcnn_benchmark.utils.model_serialization import load_state_dict
23 from pythia.utils.general import download_file
24 
25 
26 class FeatureExtractor:
27     MODEL_URL = (
28         "https://dl.fbaipublicfiles.com/pythia/detectron_model/detectron_model.pth"
29     )
30     CONFIG_URL = (
31         "https://dl.fbaipublicfiles.com/pythia/detectron_model/detectron_model.yaml"
32     )
33     MAX_SIZE = 1333
34     MIN_SIZE = 800
35 
36     def __init__(self):
37         self.args = self.get_parser().parse_args()
38         self.detection_model = self._build_detection_model()
39 
40         os.makedirs(self.args.output_folder, exist_ok=True)
41 
42     def _try_downloading_necessities(self):
43         if self.args.model_file is None:
44             print("Downloading model and configuration")
45             self.args.model_file = self.MODEL_URL.split("/")[-1]
46             self.args.config_file = self.CONFIG_URL.split("/")[-1]
47             download_file(self.MODEL_URL)
48             download_file(self.CONFIG_URL)
49 
50     def get_parser(self):
51         parser = argparse.ArgumentParser()
52         parser.add_argument(
53             "--model_file", default=None, type=str, help="Detectron model file"
54         )
55         parser.add_argument(
56             "--config_file", default=None, type=str, help="Detectron config file"
57         )
58         parser.add_argument("--batch_size", type=int, default=2, help="Batch size")
59         parser.add_argument(
60             "--num_features", type=int, default=100, help="Number of features to extract."
61         )
62         parser.add_argument(
63             "--output_folder", type=str, default="./output", help="Output folder"
64         )
65         parser.add_argument("--image_dir", type=str, help="Image directory or file")
66         parser.add_argument(
67             "--feature_name", type=str, help="The name of the feature to extract",
68             default="fc6",
69         )
70         parser.add_argument(
71             "--confidence_threshold", type=float, default=0,
72             help="Threshold of detection confidence above which boxes will be selected"
73         )
74         parser.add_argument(
75             "--background", action="store_true",
76             help="The model will output predictions for the background class when set"
77         )
78         return parser
79 
80     def _build_detection_model(self):
81         cfg.merge_from_file(self.args.config_file)
82         cfg.freeze()
83 
84         model = build_detection_model(cfg)
85         checkpoint = torch.load(self.args.model_file, map_location=torch.device("cpu"))
86 
87         load_state_dict(model, checkpoint.pop("model"))
88 
89         model.to("cuda")
90         model.eval()
91         return model
92 
93     def _image_transform(self, path):
94         img = Image.open(path)
95         im = np.array(img).astype(np.float32)
96         # IndexError: too many indices for array, grayscale images
97         if len(im.shape) < 3:
98             im = np.repeat(im[:, :, np.newaxis], 3, axis=2)
99         im = im[:, :, ::-1]
100         im -= np.array([102.9801, 115.9465, 122.7717])
101         im_shape = im.shape
102         im_height = im_shape[0]
103         im_width = im_shape[1]
104         im_size_min = np.min(im_shape[0:2])
105         im_size_max = np.max(im_shape[0:2])
106 
107         # Scale based on minimum size
108         im_scale = self.MIN_SIZE / im_size_min
109 
110         # Prevent the biggest axis from being more than max_size
111         # If bigger, scale it down
112         if np.round(im_scale * im_size_max) > self.MAX_SIZE:
113             im_scale = self.MAX_SIZE / im_size_max
114 
115         im = cv2.resize(
116             im, None, None, fx=im_scale, fy=im_scale, interpolation=cv2.INTER_LINEAR
117         )
118         img = torch.from_numpy(im).permute(2, 0, 1)
119 
120         im_info = {
121             "width": im_width,
122             "height": im_height
123         }
124 
125         return img, im_scale, im_info
126 
127     def _process_feature_extraction(
128         self, output, im_scales, im_infos, feature_name="fc6", conf_thresh=0
129     ):
130         batch_size = len(output[0]["proposals"])
131         n_boxes_per_image = [len(boxes) for boxes in output[0]["proposals"]]
132         score_list = output[0]["scores"].split(n_boxes_per_image)
133         score_list = [torch.nn.functional.softmax(x, -1) for x in score_list]
134         feats = output[0][feature_name].split(n_boxes_per_image)
135         cur_device = score_list[0].device
136 
137         feat_list = []
138         info_list = []
139 
140         for i in range(batch_size):
141             dets = output[0]["proposals"][i].bbox / im_scales[i]
142             scores = score_list[i]
143             max_conf = torch.zeros((scores.shape[0])).to(cur_device)
144             conf_thresh_tensor = torch.full_like(max_conf, conf_thresh)
145             start_index = 1
146             # Column 0 of the scores matrix is for the background class
147             if self.args.background:
148                 start_index = 0
149             for cls_ind in range(start_index, scores.shape[1]):
150                 cls_scores = scores[:, cls_ind]
151                 keep = nms(dets, cls_scores, 0.5)
152                 max_conf[keep] = torch.where(
153                     # Better than max one till now and minimally greater than conf_thresh
154                     (cls_scores[keep] > max_conf[keep]) &
155                     (cls_scores[keep] > conf_thresh_tensor[keep]),
156                     cls_scores[keep], max_conf[keep]
157                 )
158 
159             sorted_scores, sorted_indices = torch.sort(max_conf, descending=True)
160             num_boxes = (sorted_scores[:self.args.num_features] != 0).sum()
161             keep_boxes = sorted_indices[:self.args.num_features]
162             feat_list.append(feats[i][keep_boxes])
163             bbox = output[0]["proposals"][i][keep_boxes].bbox / im_scales[i]
164             # Predict the class label using the scores
165             objects = torch.argmax(scores[keep_boxes][start_index:], dim=1)
166 
167             info_list.append(
168                 {
169                     "bbox": bbox.cpu().numpy(),
170                     "num_boxes": num_boxes.item(),
171                     "objects": objects.cpu().numpy(),
172                     "image_width": im_infos[i]["width"],
173                     "image_height": im_infos[i]["height"],
174                 }
175             )
176 
177         return feat_list, info_list
178 
179     def get_detectron_features(self, image_paths):
180         img_tensor, im_scales, im_infos = [], [], []
181 
182         for image_path in image_paths:
183             im, im_scale, im_info = self._image_transform(image_path)
184             img_tensor.append(im)
185             im_scales.append(im_scale)
186             im_infos.append(im_info)
187 
188         # Image dimensions should be divisible by 32, to allow convolutions
189         # in detector to work
190         current_img_list = to_image_list(img_tensor, size_divisible=32)
191         current_img_list = current_img_list.to("cuda")
192 
193         with torch.no_grad():
194             output = self.detection_model(current_img_list)
195 
196         feat_list = self._process_feature_extraction(
197             output, im_scales, im_infos, self.args.feature_name,
198             self.args.confidence_threshold
199         )
200 
201         return feat_list
202 
203     def _chunks(self, array, chunk_size):
204         for i in range(0, len(array), chunk_size):
205             yield array[i : i + chunk_size]
206 
207     def _save_feature(self, file_name, feature, info):
208         file_base_name = os.path.basename(file_name)
209         file_base_name = file_base_name.split(".")[0]
210         info_file_base_name = file_base_name + "_info.npy"
211         file_base_name = file_base_name + ".npy"
212 
213         np.save(
214             os.path.join(self.args.output_folder, file_base_name), feature.cpu().numpy()
215         )
216         np.save(os.path.join(self.args.output_folder, info_file_base_name), info)
217 
218     def extract_features(self):
219         image_dir = self.args.image_dir
220 
221         if os.path.isfile(image_dir):
222             features, infos = self.get_detectron_features([image_dir])
223             self._save_feature(image_dir, features[0], infos[0])
224         else:
225             files = glob.glob(os.path.join(image_dir, "*.jpg"))
226             for chunk in self._chunks(files, self.args.batch_size):
227                 features, infos = self.get_detectron_features(chunk)
228                 for idx, file_name in enumerate(chunk):
229                     self._save_feature(file_name, features[idx], infos[idx])
230 
231 
232 if __name__ == "__main__":
233     feature_extractor = FeatureExtractor()
234     feature_extractor.extract_features()
235
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pythia/scripts/features/extract_features_vmb.py b/pythia/scripts/features/extract_features_vmb.py
--- a/pythia/scripts/features/extract_features_vmb.py
+++ b/pythia/scripts/features/extract_features_vmb.py
@@ -162,7 +162,7 @@
             feat_list.append(feats[i][keep_boxes])
             bbox = output[0]["proposals"][i][keep_boxes].bbox / im_scales[i]
             # Predict the class label using the scores
-            objects = torch.argmax(scores[keep_boxes][start_index:], dim=1)
+            objects = torch.argmax(scores[keep_boxes][:, start_index:], dim=1)
 
             info_list.append(
                 {

{"golden_diff": "diff --git a/pythia/scripts/features/extract_features_vmb.py b/pythia/scripts/features/extract_features_vmb.py\n--- a/pythia/scripts/features/extract_features_vmb.py\n+++ b/pythia/scripts/features/extract_features_vmb.py\n@@ -162,7 +162,7 @@\n             feat_list.append(feats[i][keep_boxes])\n             bbox = output[0][\"proposals\"][i][keep_boxes].bbox / im_scales[i]\n             # Predict the class label using the scores\n-            objects = torch.argmax(scores[keep_boxes][start_index:], dim=1)\n+            objects = torch.argmax(scores[keep_boxes][:, start_index:], dim=1)\n \n             info_list.append(\n                 {\n", "issue": "Minor bug in object detections in extract_features_vmb.py\n## \ud83d\udc1b Bug\r\n\r\n<!-- A clear and concise description of what the bug is. -->\r\n\r\nI believe this was introduced by the #127, but there's a slight numpy indexing issue that causes the objects to still be detected incorrectly. \r\n\r\nThe line in question is [here](https://github.com/facebookresearch/pythia/blob/12f67cd4f67499814bb0b3665ff14dd635800f63/pythia/scripts/features/extract_features_vmb.py#L165). It currently reads\r\n\r\n```python\r\nobjects = torch.argmax(scores[keep_boxes][start_index:], dim=1)\r\n```\r\n\r\nHowever, `scores` is a tensor of `[num_objects,  object_classes]`, so `start_index` should be indexing the second dimension. The updated line should be\r\n\r\n```python\r\nobjects = torch.argmax(scores[keep_boxes][:, start_index:], dim=1)\r\n```\r\n\r\nI can submit a pull request.\n", "before_files": [{"content": "# Requires vqa-maskrcnn-benchmark to be built and installed\n# Category mapping for visual genome can be downloaded from\n# https://dl.fbaipublicfiles.com/pythia/data/visual_genome_categories.json\n# When the --background flag is set, the index saved with key \"objects\" in\n# info_list will be +1 of the Visual Genome category mapping above and 0\n# is the background class. When the --background flag is not set, the\n# index saved with key \"objects\" in info list will match the Visual Genome\n# category mapping.\nimport argparse\nimport glob\nimport os\n\nimport cv2\nimport numpy as np\nimport torch\nfrom PIL import Image\n\nfrom maskrcnn_benchmark.config import cfg\nfrom maskrcnn_benchmark.layers import nms\nfrom maskrcnn_benchmark.modeling.detector import build_detection_model\nfrom maskrcnn_benchmark.structures.image_list import to_image_list\nfrom maskrcnn_benchmark.utils.model_serialization import load_state_dict\nfrom pythia.utils.general import download_file\n\n\nclass FeatureExtractor:\n    MODEL_URL = (\n        \"https://dl.fbaipublicfiles.com/pythia/detectron_model/detectron_model.pth\"\n    )\n    CONFIG_URL = (\n        \"https://dl.fbaipublicfiles.com/pythia/detectron_model/detectron_model.yaml\"\n    )\n    MAX_SIZE = 1333\n    MIN_SIZE = 800\n\n    def __init__(self):\n        self.args = self.get_parser().parse_args()\n        self.detection_model = self._build_detection_model()\n\n        os.makedirs(self.args.output_folder, exist_ok=True)\n\n    def _try_downloading_necessities(self):\n        if self.args.model_file is None:\n            print(\"Downloading model and configuration\")\n            self.args.model_file = self.MODEL_URL.split(\"/\")[-1]\n            self.args.config_file = self.CONFIG_URL.split(\"/\")[-1]\n            download_file(self.MODEL_URL)\n            download_file(self.CONFIG_URL)\n\n    def get_parser(self):\n        parser = argparse.ArgumentParser()\n        parser.add_argument(\n            \"--model_file\", default=None, type=str, help=\"Detectron model file\"\n        )\n        parser.add_argument(\n            \"--config_file\", default=None, type=str, help=\"Detectron config file\"\n        )\n        parser.add_argument(\"--batch_size\", type=int, default=2, help=\"Batch size\")\n        parser.add_argument(\n            \"--num_features\", type=int, default=100, help=\"Number of features to extract.\"\n        )\n        parser.add_argument(\n            \"--output_folder\", type=str, default=\"./output\", help=\"Output folder\"\n        )\n        parser.add_argument(\"--image_dir\", type=str, help=\"Image directory or file\")\n        parser.add_argument(\n            \"--feature_name\", type=str, help=\"The name of the feature to extract\",\n            default=\"fc6\",\n        )\n        parser.add_argument(\n            \"--confidence_threshold\", type=float, default=0,\n            help=\"Threshold of detection confidence above which boxes will be selected\"\n        )\n        parser.add_argument(\n            \"--background\", action=\"store_true\",\n            help=\"The model will output predictions for the background class when set\"\n        )\n        return parser\n\n    def _build_detection_model(self):\n        cfg.merge_from_file(self.args.config_file)\n        cfg.freeze()\n\n        model = build_detection_model(cfg)\n        checkpoint = torch.load(self.args.model_file, map_location=torch.device(\"cpu\"))\n\n        load_state_dict(model, checkpoint.pop(\"model\"))\n\n        model.to(\"cuda\")\n        model.eval()\n        return model\n\n    def _image_transform(self, path):\n        img = Image.open(path)\n        im = np.array(img).astype(np.float32)\n        # IndexError: too many indices for array, grayscale images\n        if len(im.shape) < 3:\n            im = np.repeat(im[:, :, np.newaxis], 3, axis=2)\n        im = im[:, :, ::-1]\n        im -= np.array([102.9801, 115.9465, 122.7717])\n        im_shape = im.shape\n        im_height = im_shape[0]\n        im_width = im_shape[1]\n        im_size_min = np.min(im_shape[0:2])\n        im_size_max = np.max(im_shape[0:2])\n\n        # Scale based on minimum size\n        im_scale = self.MIN_SIZE / im_size_min\n\n        # Prevent the biggest axis from being more than max_size\n        # If bigger, scale it down\n        if np.round(im_scale * im_size_max) > self.MAX_SIZE:\n            im_scale = self.MAX_SIZE / im_size_max\n\n        im = cv2.resize(\n            im, None, None, fx=im_scale, fy=im_scale, interpolation=cv2.INTER_LINEAR\n        )\n        img = torch.from_numpy(im).permute(2, 0, 1)\n\n        im_info = {\n            \"width\": im_width,\n            \"height\": im_height\n        }\n\n        return img, im_scale, im_info\n\n    def _process_feature_extraction(\n        self, output, im_scales, im_infos, feature_name=\"fc6\", conf_thresh=0\n    ):\n        batch_size = len(output[0][\"proposals\"])\n        n_boxes_per_image = [len(boxes) for boxes in output[0][\"proposals\"]]\n        score_list = output[0][\"scores\"].split(n_boxes_per_image)\n        score_list = [torch.nn.functional.softmax(x, -1) for x in score_list]\n        feats = output[0][feature_name].split(n_boxes_per_image)\n        cur_device = score_list[0].device\n\n        feat_list = []\n        info_list = []\n\n        for i in range(batch_size):\n            dets = output[0][\"proposals\"][i].bbox / im_scales[i]\n            scores = score_list[i]\n            max_conf = torch.zeros((scores.shape[0])).to(cur_device)\n            conf_thresh_tensor = torch.full_like(max_conf, conf_thresh)\n            start_index = 1\n            # Column 0 of the scores matrix is for the background class\n            if self.args.background:\n                start_index = 0\n            for cls_ind in range(start_index, scores.shape[1]):\n                cls_scores = scores[:, cls_ind]\n                keep = nms(dets, cls_scores, 0.5)\n                max_conf[keep] = torch.where(\n                    # Better than max one till now and minimally greater than conf_thresh\n                    (cls_scores[keep] > max_conf[keep]) &\n                    (cls_scores[keep] > conf_thresh_tensor[keep]),\n                    cls_scores[keep], max_conf[keep]\n                )\n\n            sorted_scores, sorted_indices = torch.sort(max_conf, descending=True)\n            num_boxes = (sorted_scores[:self.args.num_features] != 0).sum()\n            keep_boxes = sorted_indices[:self.args.num_features]\n            feat_list.append(feats[i][keep_boxes])\n            bbox = output[0][\"proposals\"][i][keep_boxes].bbox / im_scales[i]\n            # Predict the class label using the scores\n            objects = torch.argmax(scores[keep_boxes][start_index:], dim=1)\n\n            info_list.append(\n                {\n                    \"bbox\": bbox.cpu().numpy(),\n                    \"num_boxes\": num_boxes.item(),\n                    \"objects\": objects.cpu().numpy(),\n                    \"image_width\": im_infos[i][\"width\"],\n                    \"image_height\": im_infos[i][\"height\"],\n                }\n            )\n\n        return feat_list, info_list\n\n    def get_detectron_features(self, image_paths):\n        img_tensor, im_scales, im_infos = [], [], []\n\n        for image_path in image_paths:\n            im, im_scale, im_info = self._image_transform(image_path)\n            img_tensor.append(im)\n            im_scales.append(im_scale)\n            im_infos.append(im_info)\n\n        # Image dimensions should be divisible by 32, to allow convolutions\n        # in detector to work\n        current_img_list = to_image_list(img_tensor, size_divisible=32)\n        current_img_list = current_img_list.to(\"cuda\")\n\n        with torch.no_grad():\n            output = self.detection_model(current_img_list)\n\n        feat_list = self._process_feature_extraction(\n            output, im_scales, im_infos, self.args.feature_name,\n            self.args.confidence_threshold\n        )\n\n        return feat_list\n\n    def _chunks(self, array, chunk_size):\n        for i in range(0, len(array), chunk_size):\n            yield array[i : i + chunk_size]\n\n    def _save_feature(self, file_name, feature, info):\n        file_base_name = os.path.basename(file_name)\n        file_base_name = file_base_name.split(\".\")[0]\n        info_file_base_name = file_base_name + \"_info.npy\"\n        file_base_name = file_base_name + \".npy\"\n\n        np.save(\n            os.path.join(self.args.output_folder, file_base_name), feature.cpu().numpy()\n        )\n        np.save(os.path.join(self.args.output_folder, info_file_base_name), info)\n\n    def extract_features(self):\n        image_dir = self.args.image_dir\n\n        if os.path.isfile(image_dir):\n            features, infos = self.get_detectron_features([image_dir])\n            self._save_feature(image_dir, features[0], infos[0])\n        else:\n            files = glob.glob(os.path.join(image_dir, \"*.jpg\"))\n            for chunk in self._chunks(files, self.args.batch_size):\n                features, infos = self.get_detectron_features(chunk)\n                for idx, file_name in enumerate(chunk):\n                    self._save_feature(file_name, features[idx], infos[idx])\n\n\nif __name__ == \"__main__\":\n    feature_extractor = FeatureExtractor()\n    feature_extractor.extract_features()\n", "path": "pythia/scripts/features/extract_features_vmb.py"}], "after_files": [{"content": "# Requires vqa-maskrcnn-benchmark to be built and installed\n# Category mapping for visual genome can be downloaded from\n# https://dl.fbaipublicfiles.com/pythia/data/visual_genome_categories.json\n# When the --background flag is set, the index saved with key \"objects\" in\n# info_list will be +1 of the Visual Genome category mapping above and 0\n# is the background class. When the --background flag is not set, the\n# index saved with key \"objects\" in info list will match the Visual Genome\n# category mapping.\nimport argparse\nimport glob\nimport os\n\nimport cv2\nimport numpy as np\nimport torch\nfrom PIL import Image\n\nfrom maskrcnn_benchmark.config import cfg\nfrom maskrcnn_benchmark.layers import nms\nfrom maskrcnn_benchmark.modeling.detector import build_detection_model\nfrom maskrcnn_benchmark.structures.image_list import to_image_list\nfrom maskrcnn_benchmark.utils.model_serialization import load_state_dict\nfrom pythia.utils.general import download_file\n\n\nclass FeatureExtractor:\n    MODEL_URL = (\n        \"https://dl.fbaipublicfiles.com/pythia/detectron_model/detectron_model.pth\"\n    )\n    CONFIG_URL = (\n        \"https://dl.fbaipublicfiles.com/pythia/detectron_model/detectron_model.yaml\"\n    )\n    MAX_SIZE = 1333\n    MIN_SIZE = 800\n\n    def __init__(self):\n        self.args = self.get_parser().parse_args()\n        self.detection_model = self._build_detection_model()\n\n        os.makedirs(self.args.output_folder, exist_ok=True)\n\n    def _try_downloading_necessities(self):\n        if self.args.model_file is None:\n            print(\"Downloading model and configuration\")\n            self.args.model_file = self.MODEL_URL.split(\"/\")[-1]\n            self.args.config_file = self.CONFIG_URL.split(\"/\")[-1]\n            download_file(self.MODEL_URL)\n            download_file(self.CONFIG_URL)\n\n    def get_parser(self):\n        parser = argparse.ArgumentParser()\n        parser.add_argument(\n            \"--model_file\", default=None, type=str, help=\"Detectron model file\"\n        )\n        parser.add_argument(\n            \"--config_file\", default=None, type=str, help=\"Detectron config file\"\n        )\n        parser.add_argument(\"--batch_size\", type=int, default=2, help=\"Batch size\")\n        parser.add_argument(\n            \"--num_features\", type=int, default=100, help=\"Number of features to extract.\"\n        )\n        parser.add_argument(\n            \"--output_folder\", type=str, default=\"./output\", help=\"Output folder\"\n        )\n        parser.add_argument(\"--image_dir\", type=str, help=\"Image directory or file\")\n        parser.add_argument(\n            \"--feature_name\", type=str, help=\"The name of the feature to extract\",\n            default=\"fc6\",\n        )\n        parser.add_argument(\n            \"--confidence_threshold\", type=float, default=0,\n            help=\"Threshold of detection confidence above which boxes will be selected\"\n        )\n        parser.add_argument(\n            \"--background\", action=\"store_true\",\n            help=\"The model will output predictions for the background class when set\"\n        )\n        return parser\n\n    def _build_detection_model(self):\n        cfg.merge_from_file(self.args.config_file)\n        cfg.freeze()\n\n        model = build_detection_model(cfg)\n        checkpoint = torch.load(self.args.model_file, map_location=torch.device(\"cpu\"))\n\n        load_state_dict(model, checkpoint.pop(\"model\"))\n\n        model.to(\"cuda\")\n        model.eval()\n        return model\n\n    def _image_transform(self, path):\n        img = Image.open(path)\n        im = np.array(img).astype(np.float32)\n        # IndexError: too many indices for array, grayscale images\n        if len(im.shape) < 3:\n            im = np.repeat(im[:, :, np.newaxis], 3, axis=2)\n        im = im[:, :, ::-1]\n        im -= np.array([102.9801, 115.9465, 122.7717])\n        im_shape = im.shape\n        im_height = im_shape[0]\n        im_width = im_shape[1]\n        im_size_min = np.min(im_shape[0:2])\n        im_size_max = np.max(im_shape[0:2])\n\n        # Scale based on minimum size\n        im_scale = self.MIN_SIZE / im_size_min\n\n        # Prevent the biggest axis from being more than max_size\n        # If bigger, scale it down\n        if np.round(im_scale * im_size_max) > self.MAX_SIZE:\n            im_scale = self.MAX_SIZE / im_size_max\n\n        im = cv2.resize(\n            im, None, None, fx=im_scale, fy=im_scale, interpolation=cv2.INTER_LINEAR\n        )\n        img = torch.from_numpy(im).permute(2, 0, 1)\n\n        im_info = {\n            \"width\": im_width,\n            \"height\": im_height\n        }\n\n        return img, im_scale, im_info\n\n    def _process_feature_extraction(\n        self, output, im_scales, im_infos, feature_name=\"fc6\", conf_thresh=0\n    ):\n        batch_size = len(output[0][\"proposals\"])\n        n_boxes_per_image = [len(boxes) for boxes in output[0][\"proposals\"]]\n        score_list = output[0][\"scores\"].split(n_boxes_per_image)\n        score_list = [torch.nn.functional.softmax(x, -1) for x in score_list]\n        feats = output[0][feature_name].split(n_boxes_per_image)\n        cur_device = score_list[0].device\n\n        feat_list = []\n        info_list = []\n\n        for i in range(batch_size):\n            dets = output[0][\"proposals\"][i].bbox / im_scales[i]\n            scores = score_list[i]\n            max_conf = torch.zeros((scores.shape[0])).to(cur_device)\n            conf_thresh_tensor = torch.full_like(max_conf, conf_thresh)\n            start_index = 1\n            # Column 0 of the scores matrix is for the background class\n            if self.args.background:\n                start_index = 0\n            for cls_ind in range(start_index, scores.shape[1]):\n                cls_scores = scores[:, cls_ind]\n                keep = nms(dets, cls_scores, 0.5)\n                max_conf[keep] = torch.where(\n                    # Better than max one till now and minimally greater than conf_thresh\n                    (cls_scores[keep] > max_conf[keep]) &\n                    (cls_scores[keep] > conf_thresh_tensor[keep]),\n                    cls_scores[keep], max_conf[keep]\n                )\n\n            sorted_scores, sorted_indices = torch.sort(max_conf, descending=True)\n            num_boxes = (sorted_scores[:self.args.num_features] != 0).sum()\n            keep_boxes = sorted_indices[:self.args.num_features]\n            feat_list.append(feats[i][keep_boxes])\n            bbox = output[0][\"proposals\"][i][keep_boxes].bbox / im_scales[i]\n            # Predict the class label using the scores\n            objects = torch.argmax(scores[keep_boxes][:, start_index:], dim=1)\n\n            info_list.append(\n                {\n                    \"bbox\": bbox.cpu().numpy(),\n                    \"num_boxes\": num_boxes.item(),\n                    \"objects\": objects.cpu().numpy(),\n                    \"image_width\": im_infos[i][\"width\"],\n                    \"image_height\": im_infos[i][\"height\"],\n                }\n            )\n\n        return feat_list, info_list\n\n    def get_detectron_features(self, image_paths):\n        img_tensor, im_scales, im_infos = [], [], []\n\n        for image_path in image_paths:\n            im, im_scale, im_info = self._image_transform(image_path)\n            img_tensor.append(im)\n            im_scales.append(im_scale)\n            im_infos.append(im_info)\n\n        # Image dimensions should be divisible by 32, to allow convolutions\n        # in detector to work\n        current_img_list = to_image_list(img_tensor, size_divisible=32)\n        current_img_list = current_img_list.to(\"cuda\")\n\n        with torch.no_grad():\n            output = self.detection_model(current_img_list)\n\n        feat_list = self._process_feature_extraction(\n            output, im_scales, im_infos, self.args.feature_name,\n            self.args.confidence_threshold\n        )\n\n        return feat_list\n\n    def _chunks(self, array, chunk_size):\n        for i in range(0, len(array), chunk_size):\n            yield array[i : i + chunk_size]\n\n    def _save_feature(self, file_name, feature, info):\n        file_base_name = os.path.basename(file_name)\n        file_base_name = file_base_name.split(\".\")[0]\n        info_file_base_name = file_base_name + \"_info.npy\"\n        file_base_name = file_base_name + \".npy\"\n\n        np.save(\n            os.path.join(self.args.output_folder, file_base_name), feature.cpu().numpy()\n        )\n        np.save(os.path.join(self.args.output_folder, info_file_base_name), info)\n\n    def extract_features(self):\n        image_dir = self.args.image_dir\n\n        if os.path.isfile(image_dir):\n            features, infos = self.get_detectron_features([image_dir])\n            self._save_feature(image_dir, features[0], infos[0])\n        else:\n            files = glob.glob(os.path.join(image_dir, \"*.jpg\"))\n            for chunk in self._chunks(files, self.args.batch_size):\n                features, infos = self.get_detectron_features(chunk)\n                for idx, file_name in enumerate(chunk):\n                    self._save_feature(file_name, features[idx], infos[idx])\n\n\nif __name__ == \"__main__\":\n    feature_extractor = FeatureExtractor()\n    feature_extractor.extract_features()\n", "path": "pythia/scripts/features/extract_features_vmb.py"}]}