()
| 163 | |
| 164 | |
| 165 | def main(): |
| 166 | parser = argparse.ArgumentParser() |
| 167 | parser.add_argument("--checkpoint", required=True) |
| 168 | parser.add_argument("--config", required=True) |
| 169 | parser.add_argument("--image", required=True) |
| 170 | parser.add_argument("--stage", default="all", |
| 171 | help="all, preprocess, backbone, fpn, features, pvs") |
| 172 | parser.add_argument("--point-x", type=float, default=None) |
| 173 | parser.add_argument("--point-y", type=float, default=None) |
| 174 | args = parser.parse_args() |
| 175 | |
| 176 | os.makedirs(DUMP_DIR, exist_ok=True) |
| 177 | print(f"Dump directory: {DUMP_DIR}") |
| 178 | |
| 179 | # Load model |
| 180 | print(f"\nLoading model: {args.config}") |
| 181 | model = build_sam2( |
| 182 | config_file=args.config, |
| 183 | ckpt_path=args.checkpoint, |
| 184 | device="cpu", |
| 185 | ) |
| 186 | model = model.float() |
| 187 | model.eval() |
| 188 | |
| 189 | predictor = SAM2ImagePredictor(model) |
| 190 | |
| 191 | # Load image |
| 192 | image = np.array(Image.open(args.image).convert("RGB")) |
| 193 | print(f"Image: {args.image} ({image.shape[1]}x{image.shape[0]})") |
| 194 | |
| 195 | # Default point at center of image |
| 196 | px = args.point_x if args.point_x is not None else image.shape[1] / 2 |
| 197 | py = args.point_y if args.point_y is not None else image.shape[0] / 2 |
| 198 | point_coords = np.array([[px, py]], dtype=np.float32) |
| 199 | point_labels = np.array([1], dtype=np.int32) |
| 200 | print(f"Test point: ({px}, {py})") |
| 201 | |
| 202 | stages = args.stage.split(",") if args.stage != "all" else [ |
| 203 | "preprocess", "backbone", "fpn", "features", "pvs" |
| 204 | ] |
| 205 | |
| 206 | img_batch = None |
| 207 | |
| 208 | if "preprocess" in stages: |
| 209 | img_batch = dump_preprocessing(predictor, image) |
| 210 | |
| 211 | if "backbone" in stages: |
| 212 | if img_batch is None: |
| 213 | img_batch = dump_preprocessing(predictor, image) |
| 214 | dump_backbone(predictor, img_batch) |
| 215 | |
| 216 | if "fpn" in stages: |
| 217 | if img_batch is None: |
| 218 | img_batch = dump_preprocessing(predictor, image) |
| 219 | dump_fpn_neck(predictor, img_batch) |
| 220 | |
| 221 | if "features" in stages: |
| 222 | dump_image_features(predictor, image) |
no test coverage detected