inference batch images

Author: jinlinyi

README.md

@@ -106,6 +106,9 @@ pf_model = PerspectiveFields(version).eval().cuda()
 img_bgr = cv2.imread('assets/imgs/cityscape.jpg')
 # inference
 predictions = pf_model.inference(img_bgr=img_bgr)
+
+# alternatively, inference a batch of images
+predictions = pf_model.inference_batch(img_bgr_list=[img_bgr_0, img_bgr_1, img_bgr_2])
 ```
 - Or checkout [Live Demo 🤗](https://huggingface.co/spaces/jinlinyi/PerspectiveFields). 
 - Notebook to [Predict Perspective Fields](./notebooks/predict_perspective_fields.ipynb). 

demo/demo.py

@@ -158,4 +158,8 @@ def resize_fix_aspect_ratio(img, field, target_width=None, target_height=None):
 pitch: 48.88
 vfov: 52.82
 cx: 0.00
-cy: 0.00""")
+cy: 0.00""")
+
+print("Alternatively, inference a batch of images")
+predictions = pf_model.inference_batch(img_bgr_list=[img_bgr, img_bgr, img_bgr])
+breakpoint()

perspective2d/perspectivefields.py

@@ -204,6 +204,22 @@ def inference(self, img_bgr):
         predictions = self.forward([inputs])[0]
         return predictions
 
+    @torch.no_grad()
+    def inference_batch(self, img_bgr_list):
+        input_list = []
+        for img_bgr in img_bgr_list:
+            original_image = img_bgr.copy()
+            if self.input_format == "RGB":
+                # whether the model expects BGR inputs or RGB
+                original_image = original_image[:, :, ::-1]
+            height, width = original_image.shape[:2]
+            image = self.aug.apply_image(original_image)
+            image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
+            inputs = {"image": image, "height": height, "width": width}
+            input_list.append(inputs)
+        predictions = self.forward(input_list)
+        return predictions
+
     def forward(self, batched_inputs) -> dict:
         """
         Forward pass of the PerspectiveFields model.
@@ -249,5 +265,8 @@ def forward(self, batched_inputs) -> dict:
                 param["pred_rel_cx"] = torch.zeros_like(param["pred_vfov"])
             if "pred_rel_cy" not in param.keys():
                 param["pred_rel_cy"] = torch.zeros_like(param["pred_vfov"])
-            processed_results[0].update(param)
+            assert len(processed_results) == len(param["pred_vfov"])
+            for i in range(len(processed_results)):
+                param_tmp = {k: v[i] for k, v in param.items()}
+                processed_results[i].update(param_tmp)
         return processed_results