fix(affine): protect use_affine and support convert

Author: HowcanoeWang

src/easyidp/geotiff.py

@@ -299,6 +299,194 @@ def set_mask_polygon(
         self._mask_polygon_is_geo = is_geo
         # Clear cached binary mask when polygon changes
         self._mask = None
+    
+    def convert_to_affine(self) -> 'GeoTiff':
+        """Convert from standard storage to affine rotation storage.
+        
+        Returns a NEW GeoTiff object with imarray aligned to the mask polygon
+        rectangle, and transform including rotation. The original object is
+        not modified.
+        
+        Requires mask_polygon to be a valid rectangle (4 vertices, 90° angles).
+        
+        Returns
+        -------
+        GeoTiff
+            New GeoTiff object in affine mode. Returns self if already affine.
+        
+        Raises
+        ------
+        ValueError
+            If no mask polygon is set or polygon is not a valid rectangle.
+        
+        Example
+        -------
+        .. code-block:: python
+        
+            >>> gtiff = idp.GeoTiff('input.tif')
+            >>> gtiff.set_mask_polygon(rect_coords, is_geo=True)
+            >>> affine_gtiff = gtiff.convert_to_affine()
+            >>> affine_gtiff.use_affine
+            True
+            >>> gtiff.use_affine  # Original unchanged
+            False
+        """
+        if self._use_affine:
+            logger.warning("Already in affine mode, returning self")
+            return self
+        
+        if self._mask_polygon is None:
+            raise ValueError("No mask polygon set. Use set_mask_polygon first.")
+        
+        polygon_geo = self.mask_polygon_geo
+        is_rect, angle, bounds = self._is_valid_rectangle(polygon_geo)
+        
+        if not is_rect:
+            raise ValueError(
+                "Polygon is not a valid rectangle. Cannot convert to affine mode."
+            )
+        
+        # Perform the transformation
+        profile = self.header['profile'].copy()
+        new_imarray, new_profile = self._prepare_affine_storage(
+            self._imarray.copy(), profile, polygon_geo, angle, bounds
+        )
+        
+        # Build new header
+        new_header = self.header.copy()
+        new_header['profile'] = new_profile
+        new_header['width'] = new_profile['width']
+        new_header['height'] = new_profile['height']
+        new_header['transform'] = new_profile['transform']
+        new_header['scale'] = [abs(new_profile['transform'].a), 
+                               abs(new_profile['transform'].e)]
+        
+        # Create new GeoTiff object
+        new_gtiff = GeoTiff(imarray=new_imarray, header=new_header)
+        new_gtiff._mask_polygon = self._mask_polygon.copy()
+        new_gtiff._mask_polygon_is_geo = self._mask_polygon_is_geo
+        new_gtiff._use_affine = True
+        
+        logger.info(f"Converted to affine mode with {angle:.1f}° rotation")
+        return new_gtiff
+    
+    def convert_from_affine(self, target_bounds: tuple | None = None) -> 'GeoTiff':
+        """Convert from affine rotation storage to standard storage.
+        
+        Returns a NEW GeoTiff object with the rotated imarray transformed back
+        to axis-aligned coordinates, with a mask for the valid region.
+        The original object is not modified.
+        
+        Parameters
+        ----------
+        target_bounds : tuple, optional
+            (min_x, min_y, max_x, max_y) for the output image bounds.
+            If None, uses the bounding box of mask_polygon.
+        
+        Returns
+        -------
+        GeoTiff
+            New GeoTiff object in standard mode. Returns self if already standard.
+        
+        Raises
+        ------
+        ValueError
+            If no mask polygon is available for conversion.
+        
+        Example
+        -------
+        .. code-block:: python
+        
+            >>> gtiff = idp.GeoTiff('affine_file.tif')
+            >>> gtiff.use_affine
+            True
+            >>> standard_gtiff = gtiff.convert_from_affine()
+            >>> standard_gtiff.use_affine
+            False
+            >>> gtiff.use_affine  # Original unchanged
+            True
+        """
+        if not self._use_affine:
+            logger.warning("Already in standard mode, returning self")
+            return self
+        
+        if self._mask_polygon is None:
+            raise ValueError("No mask polygon available for conversion.")
+        
+        from rasterio.transform import Affine
+        from scipy.ndimage import map_coordinates
+        
+        polygon_geo = self.mask_polygon_geo
+        old_transform = self.header['profile']['transform']
+        
+        # Calculate target bounds
+        if target_bounds is None:
+            min_x, max_x = polygon_geo[:, 0].min(), polygon_geo[:, 0].max()
+            min_y, max_y = polygon_geo[:, 1].min(), polygon_geo[:, 1].max()
+        else:
+            min_x, min_y, max_x, max_y = target_bounds
+        
+        # Use same pixel size
+        pixel_size_x = abs(old_transform.a)
+        pixel_size_y = abs(old_transform.e)
+        
+        # Calculate output dimensions
+        out_width = int(np.ceil((max_x - min_x) / pixel_size_x))
+        out_height = int(np.ceil((max_y - min_y) / pixel_size_y))
+        
+        # New axis-aligned transform
+        new_transform = Affine.translation(min_x, max_y) * Affine.scale(pixel_size_x, -pixel_size_y)
+        
+        # Create output coordinate grids
+        out_rows, out_cols = np.mgrid[0:out_height, 0:out_width]
+        
+        # Output pixels to geo coordinates (axis-aligned)
+        geo_x = min_x + out_cols * pixel_size_x
+        geo_y = max_y - out_rows * pixel_size_y
+        
+        # Geo coordinates to input (rotated) pixel coordinates
+        inv_old = ~old_transform
+        in_cols = inv_old.a * geo_x + inv_old.b * geo_y + inv_old.c
+        in_rows = inv_old.d * geo_x + inv_old.e * geo_y + inv_old.f
+        
+        # Sample from rotated image
+        imarray = self._imarray
+        ndim = len(imarray.shape)
+        if ndim == 2:
+            new_imarray = map_coordinates(
+                imarray, [in_rows, in_cols], order=1, mode='constant', cval=0
+            ).astype(imarray.dtype)
+        else:
+            bands = imarray.shape[2]
+            new_imarray = np.zeros((out_height, out_width, bands), dtype=imarray.dtype)
+            for b in range(bands):
+                new_imarray[:, :, b] = map_coordinates(
+                    imarray[:, :, b], [in_rows, in_cols],
+                    order=1, mode='constant', cval=0
+                ).astype(imarray.dtype)
+        
+        # Build new header
+        new_profile = self.header['profile'].copy()
+        new_profile['width'] = out_width
+        new_profile['height'] = out_height
+        new_profile['transform'] = new_transform
+        
+        new_header = self.header.copy()
+        new_header['profile'] = new_profile
+        new_header['width'] = out_width
+        new_header['height'] = out_height
+        new_header['transform'] = new_transform
+        new_header['scale'] = [pixel_size_x, pixel_size_y]
+        new_header['tie_point'] = [min_x, max_y]
+        
+        # Create new GeoTiff object
+        new_gtiff = GeoTiff(imarray=new_imarray, header=new_header)
+        new_gtiff._mask_polygon = self._mask_polygon.copy()
+        new_gtiff._mask_polygon_is_geo = self._mask_polygon_is_geo
+        new_gtiff._use_affine = False
+        
+        logger.info("Converted from affine to standard mode")
+        return new_gtiff
 
     def _polygon_pixel_to_geo(self, polygon: np.ndarray) -> np.ndarray:
         """Convert pixel polygon to geo coordinates using transform."""
@@ -743,12 +931,11 @@ def save(
             if use_affine:
                 is_rect, angle, bounds = self._is_valid_rectangle(polygon_geo)
                 if is_rect:
-                    # Prepare affine rotated storage
+                    # Prepare affine rotated storage (only for saved file)
                     imarray, profile = self._prepare_affine_storage(
                         imarray, profile, polygon_geo, angle, bounds
                     )
                     logger.info(f"Affine mode: saved with {angle:.1f}° rotation")
-                    self._use_affine = True
                     apply_mask = False  # No mask needed in affine mode
                 else:
                     logger.warning(
@@ -811,8 +998,8 @@ def _is_valid_rectangle(
         bounds : tuple
             (origin_x, origin_y, width, height) of the rectangle.
         """
-        # Remove closure point if present
-        pts = polygon[:-1] if np.allclose(polygon[0], polygon[-1]) else polygon
+        # Remove closure point if present (use strict atol for geo coords)
+        pts = polygon[:-1] if np.allclose(polygon[0], polygon[-1], rtol=0, atol=1e-6) else polygon
 
         if len(pts) != 4:
             return False, 0.0, ()

tests/test_geotiff.py

@@ -1058,3 +1058,177 @@ def test_one_raw_roi2geotiff_stores_polygon(self, shared_data):
         assert gtiff.mask_polygon is not None
         # Polygon should match input (approximately closed)
         assert len(gtiff.mask_polygon) >= len(roi_geo) - 1
+
+
+class TestAffineConversion:
+    """Tests for affine mode conversion functions."""
+    
+    def test_convert_to_affine_returns_new_object(self, shared_data):
+        """convert_to_affine returns a new GeoTiff, original unchanged."""
+        test_data = shared_data['test_data']
+        
+        # Load and set polygon
+        gtiff = idp.GeoTiff(test_data.pix4d.lotus_dom_part)
+        gtiff._imarray = gtiff.imarray  # Force load
+        original_shape = gtiff.imarray.shape
+        
+        # Set axis-aligned rectangle
+        polygon = np.array([
+            [368025.0, 3955479.0],
+            [368027.0, 3955479.0],
+            [368027.0, 3955477.0],
+            [368025.0, 3955477.0],
+        ])
+        gtiff.set_mask_polygon(polygon, is_geo=True)
+        
+        # Convert to affine
+        affine_gtiff = gtiff.convert_to_affine()
+        
+        # Verify new object
+        assert affine_gtiff is not gtiff
+        assert affine_gtiff.use_affine is True
+        
+        # Verify original unchanged
+        assert gtiff.use_affine is False
+        assert gtiff.imarray.shape == original_shape
+    
+    def test_convert_from_affine_returns_new_object(self, shared_data):
+        """convert_from_affine returns a new GeoTiff, original unchanged."""
+        test_data = shared_data['test_data']
+        
+        # Load, set polygon, convert to affine
+        gtiff = idp.GeoTiff(test_data.pix4d.lotus_dom_part)
+        gtiff._imarray = gtiff.imarray
+        
+        polygon = np.array([
+            [368025.0, 3955479.0],
+            [368027.0, 3955479.0],
+            [368027.0, 3955477.0],
+            [368025.0, 3955477.0],
+        ])
+        gtiff.set_mask_polygon(polygon, is_geo=True)
+        
+        affine_gtiff = gtiff.convert_to_affine()
+        affine_shape = affine_gtiff.imarray.shape
+        
+        # Convert back from affine
+        standard_gtiff = affine_gtiff.convert_from_affine()
+        
+        # Verify new object
+        assert standard_gtiff is not affine_gtiff
+        assert standard_gtiff.use_affine is False
+        
+        # Verify affine original unchanged
+        assert affine_gtiff.use_affine is True
+        assert affine_gtiff.imarray.shape == affine_shape
+    
+    def test_convert_roundtrip_data_consistency(self, shared_data):
+        """Roundtrip conversion preserves polygon and basic properties."""
+        test_data = shared_data['test_data']
+        
+        # Load original image
+        gtiff = idp.GeoTiff(test_data.pix4d.lotus_dom_part)
+        gtiff._imarray = gtiff.imarray
+        original_shape = gtiff.imarray.shape
+        
+        # Set axis-aligned rectangle
+        polygon = np.array([
+            [368025.0, 3955479.0],
+            [368027.0, 3955479.0],
+            [368027.0, 3955477.0],
+            [368025.0, 3955477.0],
+        ])
+        gtiff.set_mask_polygon(polygon, is_geo=True)
+        
+        # Convert to affine
+        affine_gtiff = gtiff.convert_to_affine()
+        assert affine_gtiff.use_affine is True
+        # Affine should be smaller (cropped to rectangle)
+        assert affine_gtiff.width <= gtiff.width
+        assert affine_gtiff.height <= gtiff.height
+        
+        # Convert back
+        roundtrip_gtiff = affine_gtiff.convert_from_affine()
+        assert roundtrip_gtiff.use_affine is False
+        
+        # Polygon should be preserved through both conversions
+        assert roundtrip_gtiff.mask_polygon is not None
+        np.testing.assert_allclose(
+            roundtrip_gtiff.mask_polygon_geo[:4], 
+            polygon, 
+            atol=1e-6
+        )
+        
+        # Roundtrip should have data (not all zeros)
+        assert roundtrip_gtiff.imarray.sum() > 0
+    
+    def test_convert_preserves_original(self, shared_data):
+        """Verify original GeoTiff is completely unchanged after conversion."""
+        test_data = shared_data['test_data']
+        
+        gtiff = idp.GeoTiff(test_data.pix4d.lotus_dom_part)
+        gtiff._imarray = gtiff.imarray
+        
+        # Store original state
+        original_imarray = gtiff.imarray.copy()
+        original_use_affine = gtiff.use_affine
+        original_width = gtiff.width
+        original_height = gtiff.height
+        
+        polygon = np.array([
+            [368025.0, 3955479.0],
+            [368027.0, 3955479.0],
+            [368027.0, 3955477.0],
+            [368025.0, 3955477.0],
+        ])
+        gtiff.set_mask_polygon(polygon, is_geo=True)
+        
+        # Convert to affine
+        _ = gtiff.convert_to_affine()
+        
+        # Verify original is unchanged
+        assert gtiff.use_affine == original_use_affine
+        assert gtiff.width == original_width
+        assert gtiff.height == original_height
+        np.testing.assert_array_equal(gtiff.imarray, original_imarray)
+    
+    def test_convert_non_rectangle_raises(self, shared_data):
+        """convert_to_affine raises error for non-rectangle polygons."""
+        test_data = shared_data['test_data']
+        
+        gtiff = idp.GeoTiff(test_data.pix4d.lotus_dom_part)
+        gtiff._imarray = gtiff.imarray
+        
+        # Set triangle (not rectangle)
+        triangle = np.array([
+            [368025.0, 3955479.0],
+            [368027.0, 3955479.0],
+            [368026.0, 3955477.0],
+        ])
+        gtiff.set_mask_polygon(triangle, is_geo=True)
+        
+        # Should raise ValueError
+        with pytest.raises(ValueError, match="not a valid rectangle"):
+            gtiff.convert_to_affine()
+    
+    def test_already_affine_returns_self(self, shared_data):
+        """convert_to_affine returns self if already in affine mode."""
+        test_data = shared_data['test_data']
+        
+        gtiff = idp.GeoTiff(test_data.pix4d.lotus_dom_part)
+        gtiff._imarray = gtiff.imarray
+        
+        polygon = np.array([
+            [368025.0, 3955479.0],
+            [368027.0, 3955479.0],
+            [368027.0, 3955477.0],
+            [368025.0, 3955477.0],
+        ])
+        gtiff.set_mask_polygon(polygon, is_geo=True)
+        
+        # Convert twice
+        affine1 = gtiff.convert_to_affine()
+        affine2 = affine1.convert_to_affine()
+        
+        # Second call should return same object
+        assert affine2 is affine1