nargs and smaller output size

stefanklut · Dec 5, 2023 · 195e057 · 195e057
1 parent fbad945
commit 195e057
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/eval.py b/eval.py
@@ -30,7 +30,7 @@ def get_arguments() -> argparse.Namespace:
     detectron2_args = parser.add_argument_group("detectron2")
 
     detectron2_args.add_argument("-c", "--config", help="config file", required=True)
-    detectron2_args.add_argument("--opts", nargs="+", help="optional args to change", default=[])
+    detectron2_args.add_argument("--opts", nargs="+", help="optional args to change", action="extend", default=[])
 
     io_args = parser.add_argument_group("IO")
     # io_args.add_argument("-t", "--train", help="Train input folder/file",
@@ -125,12 +125,16 @@ def create_pred_visualization(image_filename):
             image = load_image(image_filename)
             logger.info(f"Predict: {image_filename}")
             outputs = predictor(image)
-            pred = torch.argmax(outputs[0]["sem_seg"], dim=-3).to("cpu")
+            sem_seg = outputs[0]["sem_seg"]
+            sem_seg = torch.nn.functional.interpolate(
+                sem_seg[None], size=(image.shape[0], image.shape[1]), mode="bilinear", align_corners=False
+            )[0]
+            sem_seg = torch.argmax(sem_seg, dim=-3).cpu().numpy()
             # outputs["panoptic_seg"] = (outputs["panoptic_seg"][0].to("cpu"),
             #                            outputs["panoptic_seg"][1])
             vis_im = Visualizer(image.copy(), metadata=metadata, scale=1)
 
-            vis_im = vis_im.draw_sem_seg(pred, alpha=0.4)
+            vis_im = vis_im.draw_sem_seg(sem_seg, alpha=0.4)
             return vis_im.get_image()
 
         fig, axes = plt.subplots(1, 2)

diff --git a/run.py b/run.py
@@ -31,7 +31,7 @@ def get_arguments() -> argparse.Namespace:
     detectron2_args = parser.add_argument_group("detectron2")
 
     detectron2_args.add_argument("-c", "--config", help="config file", required=True)
-    detectron2_args.add_argument("--opts", nargs=argparse.REMAINDER, help="optional args to change", default=[])
+    detectron2_args.add_argument("--opts", nargs="+", help="optional args to change", action="extend", default=[])
 
     io_args = parser.add_argument_group("IO")
     io_args.add_argument("-i", "--input", nargs="+", help="Input folder", type=str, action="extend", required=True)
@@ -121,7 +121,7 @@ def cpu_call(self, original_image: np.ndarray):
                 # whether the model expects BGR inputs or RGB
                 image = image[[2, 1, 0], :, :]
 
-            inputs = {"image": image, "height": height, "width": width}
+            inputs = {"image": image, "height": image.shape[1], "width": image.shape[2]}
             predictions = self.model([inputs])[0]
 
         return predictions, height, width

diff --git a/xml_viewer.py b/xml_viewer.py
@@ -27,7 +27,7 @@ def get_arguments() -> argparse.Namespace:
     detectron2_args = parser.add_argument_group("detectron2")
 
     detectron2_args.add_argument("-c", "--config", help="config file", required=True)
-    detectron2_args.add_argument("--opts", nargs="+", help="optional args to change", default=[])
+    detectron2_args.add_argument("--opts", nargs="+", help="optional args to change", action="extend", default=[])
 
     io_args = parser.add_argument_group("IO")
     io_args.add_argument("-i", "--input", help="Input folder/files", nargs="+", default=[], required=True, type=str)