vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/UFM/uniflowmatch/cli.py

@@ -0,0 +1,217 @@
+#!/usr/bin/env python3
+"""
+Command-line interface for UFM.
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+
+
+def main():
+    """Main CLI entry point."""
+    parser = argparse.ArgumentParser(description="UFM: Unified Dense Correspondence with Flow", prog="ufm")
+
+    subparsers = parser.add_subparsers(dest="command", help="Available commands")
+
+    # Demo command
+    demo_parser = subparsers.add_parser("demo", help="Launch interactive Gradio demo")
+    demo_parser.add_argument("--port", type=int, default=7860, help="Port to run demo on (default: 7860)")
+    demo_parser.add_argument("--share", action="store_true", help="Create public sharing link")
+    demo_parser.add_argument(
+        "--model", choices=["base", "refine"], default="base", help="Model variant to use (default: base)"
+    )
+
+    # Inference command
+    infer_parser = subparsers.add_parser("infer", help="Run inference on image pairs")
+    infer_parser.add_argument("source", help="Source image path")
+    infer_parser.add_argument("target", help="Target image path")
+    infer_parser.add_argument("--output", "-o", help="Output directory (default: current directory)")
+    infer_parser.add_argument(
+        "--model", choices=["base", "refine"], default="base", help="Model variant to use (default: base)"
+    )
+
+    # Test command
+    test_parser = subparsers.add_parser("test", help="Test installation")
+
+    args = parser.parse_args()
+
+    if args.command == "demo":
+        launch_demo(args)
+    elif args.command == "infer":
+        run_inference(args)
+    elif args.command == "test":
+        test_installation()
+    else:
+        parser.print_help()
+
+
+def launch_demo(args):
+    """Launch the Gradio demo."""
+    try:
+        # Import here to avoid slow startup for other commands
+        from gradio_demo import create_demo, initialize_model
+
+        print(f"Launching UFM demo with {args.model} model...")
+        print(f"Demo will be available at: http://localhost:{args.port}")
+
+        # Initialize model
+        use_refinement = args.model == "refine"
+        model_loaded = initialize_model(use_refinement=use_refinement)
+
+        if not model_loaded:
+            print("Error: Failed to load model. Check installation and internet connection.")
+            sys.exit(1)
+
+        # Create and launch demo
+        demo = create_demo()
+        demo.launch(
+            share=args.share,
+            server_port=args.port,
+            server_name="127.0.0.1",
+            show_error=True,
+        )
+
+    except ImportError as e:
+        print(f"Error importing demo dependencies: {e}")
+        print("Please install demo dependencies: pip install -e '.[demo]'")
+        sys.exit(1)
+    except Exception as e:
+        print(f"Error launching demo: {e}")
+        sys.exit(1)
+
+
+def run_inference(args):
+    """Run inference on image pair."""
+    try:
+        import cv2
+        import flow_vis
+        import numpy as np
+        import torch
+
+        from uniflowmatch.models.ufm import UniFlowMatchClassificationRefinement, UniFlowMatchConfidence
+        from uniflowmatch.utils.viz import warp_image_with_flow
+
+        # Load images
+        source_img = cv2.imread(args.source)
+        target_img = cv2.imread(args.target)
+
+        if source_img is None or target_img is None:
+            print("Error: Could not load one or both images")
+            sys.exit(1)
+
+        # Convert BGR to RGB
+        source_rgb = cv2.cvtColor(source_img, cv2.COLOR_BGR2RGB)
+        target_rgb = cv2.cvtColor(target_img, cv2.COLOR_BGR2RGB)
+
+        # Load model
+        if args.model == "refine":
+            model = UniFlowMatchClassificationRefinement.from_pretrained("infinity1096/UFM-Refine")
+        else:
+            model = UniFlowMatchConfidence.from_pretrained("infinity1096/UFM-Base")
+        model.eval()
+
+        print("Running inference...")
+
+        # Run inference
+        with torch.no_grad():
+            result = model.predict_correspondences_batched(
+                source_image=torch.from_numpy(source_rgb),
+                target_image=torch.from_numpy(target_rgb),
+            )
+
+            flow = result.flow.flow_output[0].cpu().numpy()
+            covisibility = result.covisibility.mask[0].cpu().numpy()
+
+        # Save outputs
+        output_dir = Path(args.output) if args.output else Path.cwd()
+        output_dir.mkdir(exist_ok=True)
+
+        # Save flow visualization
+        flow_vis_img = flow_vis.flow_to_color(flow.transpose(1, 2, 0))
+        cv2.imwrite(str(output_dir / "flow_visualization.png"), cv2.cvtColor(flow_vis_img, cv2.COLOR_RGB2BGR))
+
+        # Save covisibility mask
+        cv2.imwrite(str(output_dir / "covisibility_mask.png"), (covisibility * 255).astype(np.uint8))
+
+        # Save warped image
+        warped_image = warp_image_with_flow(source_rgb, None, target_rgb, flow.transpose(1, 2, 0))
+        warped_image = covisibility[..., None] * warped_image + (1 - covisibility[..., None]) * 255 * np.ones_like(
+            warped_image
+        )
+        cv2.imwrite(str(output_dir / "warped_source.png"), cv2.cvtColor(warped_image, cv2.COLOR_RGB2BGR))
+
+        print(f"Results saved to: {output_dir}")
+        print("- flow_visualization.png")
+        print("- covisibility_mask.png")
+        print("- warped_source.png")
+
+    except ImportError as e:
+        print(f"Error importing dependencies: {e}")
+        print("Please ensure all dependencies are installed")
+        sys.exit(1)
+    except Exception as e:
+        print(f"Error during inference: {e}")
+        sys.exit(1)
+
+
+def test_installation():
+    """Test the installation."""
+    print("Testing UFM installation...")
+
+    try:
+        # Test imports
+        import torch
+
+        print(f"✓ PyTorch {torch.__version__}")
+
+        import torchvision
+
+        print(f"✓ TorchVision {torchvision.__version__}")
+
+        import numpy
+
+        print(f"✓ NumPy {numpy.__version__}")
+
+        import cv2
+
+        print(f"✓ OpenCV {cv2.__version__}")
+
+        # Test UFM imports
+        from uniflowmatch.models.ufm import UniFlowMatchConfidence
+
+        print("✓ UFM model imports")
+
+        # Test CUDA if available
+        if torch.cuda.is_available():
+            print(f"✓ CUDA {torch.version.cuda} available")
+            print(f"  GPU: {torch.cuda.get_device_name(0)}")
+        else:
+            print("⚠ CUDA not available (CPU-only mode)")
+
+        # Test model loading (lightweight test)
+        print("Testing model loading...")
+        try:
+            # This should work if HuggingFace hub is accessible
+            from huggingface_hub import hf_hub_download
+
+            print("✓ HuggingFace Hub accessible")
+        except Exception:
+            print("⚠ HuggingFace Hub not accessible (may affect model downloading)")
+
+        print("\n✅ Installation test completed successfully!")
+        print("Run 'ufm demo' to launch the interactive demo")
+
+    except ImportError as e:
+        print(f"❌ Import error: {e}")
+        print("Please check your installation")
+        sys.exit(1)
+    except Exception as e:
+        print(f"❌ Unexpected error: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

vismatch/third_party/UFM/uniflowmatch/models/__init__.py

@@ -0,0 +1,25 @@
+"""Models subpackage for `uniflowmatch`."""
+
+from .base import (
+    UFMClassificationRefinementOutput,
+    UFMFlowFieldOutput,
+    UFMMaskFieldOutput,
+    UFMOutputInterface,
+    UniFlowMatchModelsBase,
+)
+from .ufm import (
+    UniFlowMatch,
+    UniFlowMatchClassificationRefinement,
+    UniFlowMatchConfidence,
+)
+
+__all__ = [
+    "UFMClassificationRefinementOutput",
+    "UFMFlowFieldOutput",
+    "UFMMaskFieldOutput",
+    "UFMOutputInterface",
+    "UniFlowMatchModelsBase",
+    "UniFlowMatch",
+    "UniFlowMatchClassificationRefinement",
+    "UniFlowMatchConfidence",
+]

vismatch/third_party/UFM/uniflowmatch/models/base.py

@@ -0,0 +1,334 @@
+"""
+Base class of the UniFlowMatch training system.
+"""
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+
+
+@dataclass
+class UFMFlowFieldOutput:
+    """
+    Output interface of the flow field prediction network.
+    """
+
+    flow_output: torch.Tensor
+    flow_covariance: Optional[torch.Tensor] = None
+    flow_covariance_inv: Optional[torch.Tensor] = None
+    flow_covariance_log_det: Optional[torch.Tensor] = None
+
+
+@dataclass
+class UFMMaskFieldOutput:
+    """
+    Output interface of the mask prediction network.
+    """
+
+    mask: torch.Tensor
+    logits: torch.Tensor
+
+
+@dataclass
+class UFMClassificationRefinementOutput:
+    """
+    Output interface of the classification refinement network.
+    """
+
+    # the flow output of the regression step, with shape [B, 2, H, W].
+    # it is the initial flow output, which is used to get the first local feature maps for the residual.
+    regression_flow_output: torch.Tensor
+
+    # residual is the output of the refinement step, with shape [B, 2, H, W].
+    # it is added to the initial flow output to get the final flow output.
+    residual: torch.Tensor
+
+    # log_softmax is
+    #   the logarithm of
+    #   the softmax of
+    #   similarity of the pixel's feature
+    #       to that of its neighborhood of the flow prediction
+    #       in the other image.
+    # it have shape [B, H, W, P, P], the similarity of pixel at [b, h, w] to its neighborhood [P, P] centered at regression_flow_output[b, h, w]
+    log_softmax: torch.Tensor
+
+    feature_map_0: torch.Tensor
+    feature_map_1: torch.Tensor
+
+
+@dataclass
+class UFMOutputInterface:
+    """
+    Output interface of the UniFlowMatch training system.
+    """
+
+    flow: Optional[UFMFlowFieldOutput] = None
+
+    # Refinement output (for training and visualization)
+    classification_refinement: Optional[UFMClassificationRefinementOutput] = None
+
+    # auxiliary ouputs
+    covisibility: Optional[UFMMaskFieldOutput] = None
+
+
+from uniception.models.encoders.image_normalizations import IMAGE_NORMALIZATION_DICT
+
+from ..utils.flow_resizing import (
+    AutomaticShapeSelection,
+    ResizeToFixedManipulation,
+    unmap_predicted_channels,
+    unmap_predicted_flow,
+)
+
+
+class UniFlowMatchModelsBase(torch.nn.Module):
+    def __init__(self, inference_resolution: Optional[Union[List[Tuple[int, int]], Tuple[int, int]]] = None):
+        super().__init__()
+
+        if inference_resolution is None:
+            inference_resolution = [(560, 420)]
+
+        if isinstance(inference_resolution[0], int):  # Handle the case for single resolution
+            inference_resolution = [inference_resolution]
+
+        self.inference_resolution = inference_resolution
+
+        self.image_scaler = AutomaticShapeSelection(
+            *[ResizeToFixedManipulation((resolution[1], resolution[0])) for resolution in inference_resolution],
+            strategy="closest_aspect",  # will inference on the trained aspect ratio that is closest to the input image 1
+        )
+
+    def forward(self, view1, view2) -> UFMOutputInterface:
+        """
+        Forward interface of correspondence prediction networks.
+
+        Args:
+        - view1 (Dict[str, Any]): Input view 1
+          - img (torch.Tensor): BCHW image tensor normalized according to encoder's data_norm_type
+          - instance (List[int]): List of instance indices, or id of the input image
+          - data_norm_type (str): Data normalization type, see uniception.models.encoders.IMAGE_NORMALIZATION_DICT
+        - view2 (Dict[str, Any]): Input view 2
+          - (same structure as view1)
+        Returns:
+        - Dict[str, Any]: Output results
+          - flow [Required] (Dict[str, torch.Tensor]): Flow output
+            - [Required] flow_output (torch.Tensor): Flow output tensor, BCHW
+            - [Optional] flow_covariance
+            - [Optional] flow_covariance_inv
+            - [Optional] flow_covariance_log_det
+          - occlusion [Optional] (Dict[str, torch.Tensor]): Occlusion output
+            - [Optional] mask
+            - [Optional] logits
+        """
+        raise NotImplementedError("Implement this method in derived classes")
+
+    def get_parameter_groups(self) -> Dict[str, torch.nn.ParameterList]:
+        """
+        Get parameter groups for optimizer. This methods guides the optimizer
+        to apply correct learning rate to different parts of the model.
+
+        Returns:
+        - Dict[str, torch.nn.ParameterList]: Parameter groups for optimizer
+        """
+
+        raise NotImplementedError("Implement this method in derived classes")
+
+    def predict_correspondences_batched(
+        self,
+        source_image: torch.Tensor,
+        target_image: torch.Tensor,
+        data_norm_type: Optional[str] = None,
+    ) -> UFMOutputInterface:
+        """
+        Predict correspondences between source and target images.
+
+        This method generates random correspondences for demonstration purposes.
+
+        Args:
+            source_image (torch.Tensor): Tensor of shape BCHW/BHWC/CHW/HWC, dtype of uint8 or float32 The source image tensor.
+            target_image (torch.Tensor): Tensor of shape BCHW/BHWC/CHW/HWC, dtype of uint8 or float32 The target image tensor.
+
+        Returns:
+            UFMOutputInterface:
+                - flow
+                    - flow_output       (torch.Tensor): Tensor of shape (B, 2, H, W) representing the flow output in the original image space.
+                - occlusion
+                    - mask              (torch.Tensor): Tensor of shape (B, H, W) representing the covisibility in range [0, 1]. 1 = fully covisible, 0 = fully occluded or out of range.
+        """
+
+        assert isinstance(source_image, torch.Tensor) and isinstance(
+            target_image, torch.Tensor
+        ), "source_image and target_image must be torch.Tensors"
+        assert source_image.dim() in [3, 4], "source_image must have dimensions 3 or 4"
+        assert target_image.dim() in [3, 4], "target_image must have dimensions 3 or 4"
+
+        batched = source_image.dim() == 4
+
+        if not batched:
+            # add batch dimension
+            source_image = source_image.unsqueeze(0)
+            target_image = target_image.unsqueeze(0)
+
+        # check the channel
+        if source_image.shape[1] == 3 and target_image.shape[1] == 3:
+            pass  # do nothing because the image is in BCHW format
+        elif source_image.shape[-1] == 3 and target_image.shape[-1] == 3:
+            # convert to BCHW
+            source_image = source_image.permute(0, 3, 1, 2)
+            target_image = target_image.permute(0, 3, 1, 2)
+        else:
+            raise ValueError("source_image and target_image must have 3 channels in either BCHW or BHWC format")
+
+        required_data_norm_type = self.encoder.data_norm_type
+
+        image_device = source_image.device
+
+        if source_image.dtype == torch.float32:
+            assert data_norm_type is not None, "data_norm_type must be provided for float32 images"
+            assert (
+                data_norm_type in IMAGE_NORMALIZATION_DICT
+            ), f"data_norm_type must be one of {list(IMAGE_NORMALIZATION_DICT.keys())}"
+
+            if data_norm_type != required_data_norm_type:
+                # apply transformation to the correct from the old normalization
+                prev_mean = (
+                    IMAGE_NORMALIZATION_DICT[data_norm_type].mean.view(1, 3, 1, 1).to(image_device, non_blocking=True)
+                )
+                prev_std = (
+                    IMAGE_NORMALIZATION_DICT[data_norm_type].std.view(1, 3, 1, 1).to(image_device, non_blocking=True)
+                )
+                mean = (
+                    IMAGE_NORMALIZATION_DICT[required_data_norm_type]
+                    .mean.view(1, 3, 1, 1)
+                    .to(image_device, non_blocking=True)
+                )
+                std = (
+                    IMAGE_NORMALIZATION_DICT[required_data_norm_type]
+                    .std.view(1, 3, 1, 1)
+                    .to(image_device, non_blocking=True)
+                )
+
+                source_image = source_image * (prev_std / std) + (prev_mean - mean) / std
+                target_image = target_image * (prev_std / std) + (prev_mean - mean) / std
+
+        elif source_image.dtype == torch.uint8:
+            # convert into float32 and apply normalization
+            mean = (
+                IMAGE_NORMALIZATION_DICT[required_data_norm_type]
+                .mean.view(1, 3, 1, 1)
+                .to(image_device, non_blocking=True)
+            )
+            std = (
+                IMAGE_NORMALIZATION_DICT[required_data_norm_type]
+                .std.view(1, 3, 1, 1)
+                .to(image_device, non_blocking=True)
+            )
+
+            source_image = (source_image.float() / 255.0 - mean) / std
+            target_image = (target_image.float() / 255.0 - mean) / std
+        else:
+            raise ValueError("source_image and target_image must be of type torch.float32 or torch.uint8")
+
+        # Now all the inputs are normalized according to the model's encoder and organized in BCHW format
+        return self._predict_correspondences_batched(source_image, target_image)
+
+    def _predict_correspondences_batched(
+        self,
+        source_image: torch.Tensor,
+        target_image: torch.Tensor,
+    ) -> UFMOutputInterface:
+        assert isinstance(source_image, torch.Tensor), "source_image must be a torch.Tensor"
+        assert isinstance(target_image, torch.Tensor), "target_image must be a torch.Tensor"
+
+        assert source_image.dim() == 4, "source_image must be of shape (B, 3, H, W)"
+        assert target_image.dim() == 4, "target_image must be of shape (B, 3, H, W)"
+        assert source_image.shape[1] == 3, "source_image must be of shape (B, 3, H, W)"
+        assert target_image.shape[1] == 3, "target_image must be of shape (B, 3, H, W)"
+
+        assert source_image.dtype == torch.float32, "source_image must be of dtype torch.float32"
+        assert target_image.dtype == torch.float32, "target_image must be of dtype torch.float32"
+
+        source_shape_hw = source_image.shape[2:]
+        target_shape_hw = target_image.shape[2:]
+
+        # Scale images to one of the model's trained resolution.
+        (
+            scaled_img0,  # The scaled source image
+            scaled_img1,  # The scaled target image
+            img0_region_source,  # Where in the source image is captured in the scaled image
+            img1_region_source,  # Where in the target image is captured in the scaled image
+            img0_region_representation,  # Region in the source image is captured in this region in the scaled image
+            img1_region_representation,  # same as above, but for the target image
+        ) = self.image_scaler(source_image.permute(0, 2, 3, 1), target_image.permute(0, 2, 3, 1))
+
+        scaled_img0 = scaled_img0.permute(0, 3, 1, 2)
+        scaled_img1 = scaled_img1.permute(0, 3, 1, 2)
+
+        # Run a forward pass
+        view1 = {"img": scaled_img0, "symmetrized": False, "data_norm_type": self.encoder.data_norm_type}
+        view2 = {"img": scaled_img1, "symmetrized": False, "data_norm_type": self.encoder.data_norm_type}
+
+        with torch.no_grad():
+            with torch.autocast("cuda", torch.bfloat16):
+                result = self(view1, view2)
+
+        rescaled_ufm_result = UFMOutputInterface()
+
+        # rescale flow
+        flow_output = result.flow.flow_output
+        flow_unmapped, flow_unmap_validity = unmap_predicted_flow(
+            flow_output,
+            img0_region_representation,
+            img1_region_representation,
+            img0_region_source,
+            img1_region_source,
+            source_shape_hw,
+            target_shape_hw,
+        )
+
+        rescaled_ufm_result.flow = UFMFlowFieldOutput(
+            flow_output=flow_unmapped,
+        )
+
+        # rescale covariance if it exists
+        if result.flow.flow_covariance is not None:
+            flow_covariance = result.flow.flow_covariance
+            flow_covariance_unmapped, _ = unmap_predicted_channels(
+                flow_covariance,
+                img0_region_representation,
+                img0_region_source,
+                source_shape_hw,
+            )
+
+            # scale covariance in the correct way
+            w_pred = scaled_img0.shape[3]
+            h_pred = scaled_img0.shape[2]
+
+            w_final = source_shape_hw[1]
+            h_final = source_shape_hw[0]
+
+            w_ratio, h_ratio = w_final / w_pred, h_final / h_pred
+
+            flow_covariance_unmapped *= (
+                torch.tensor([w_ratio**2, h_ratio**2, w_ratio * h_ratio])
+                .view(1, 3, 1, 1)
+                .to(flow_covariance_unmapped.device)
+            )
+
+            rescaled_ufm_result.flow.flow_covariance = flow_covariance_unmapped
+
+        # rescale occlusion if it exists
+        if result.covisibility is not None:
+            occlusion_mask = result.covisibility.mask
+            covisibility_unmapped, _ = unmap_predicted_channels(
+                occlusion_mask,
+                img0_region_representation,
+                img0_region_source,
+                source_shape_hw,
+            )
+
+            covisibility_unmapped = covisibility_unmapped.squeeze(1)
+            rescaled_ufm_result.covisibility = UFMMaskFieldOutput(mask=covisibility_unmapped, logits=None)
+
+        return rescaled_ufm_result