Add support for GeoTIFF

Author: giswqs

samgeo/common.py

@@ -3433,3 +3433,83 @@ def on_slider_change(change):
 
     # Show the plot
     plt.show()
+
+
+def make_temp_dir(**kwargs) -> str:
+    """Create a temporary directory and return the path.
+
+    Returns:
+        str: The path to the temporary directory.
+    """
+    import tempfile
+
+    temp_dir = tempfile.mkdtemp(**kwargs)
+    return temp_dir
+
+
+def geotiff_to_jpg(geotiff_path: str, output_path: str) -> None:
+    """Convert a GeoTIFF file to a JPG file.
+
+    Args:
+        geotiff_path (str): The path to the input GeoTIFF file.
+        output_path (str): The path to the output JPG file.
+    """
+
+    from PIL import Image
+
+    # Open the GeoTIFF file
+    with rasterio.open(geotiff_path) as src:
+        # Read the first band (for grayscale) or all bands
+        array = src.read()
+
+        # If the array has more than 3 bands, reduce it to the first 3 (RGB)
+        if array.shape[0] >= 3:
+            array = array[:3, :, :]  # Select the first 3 bands (R, G, B)
+        elif array.shape[0] == 1:
+            # For single-band images, repeat the band to create a grayscale RGB
+            array = np.repeat(array, 3, axis=0)
+
+        # Transpose the array from (bands, height, width) to (height, width, bands)
+        array = np.transpose(array, (1, 2, 0))
+
+        # Normalize the array to 8-bit (0-255) range for JPG
+        array = array.astype(np.float32)
+        array -= array.min()
+        array /= array.max()
+        array *= 255
+        array = array.astype(np.uint8)
+
+        # Convert to a PIL Image and save as JPG
+        image = Image.fromarray(array)
+        image.save(output_path)
+
+
+def geotiff_to_jpg_batch(input_folder: str, output_folder: str = None) -> str:
+    """Convert all GeoTIFF files in a folder to JPG files.
+
+    Args:
+        input_folder (str): The path to the folder containing GeoTIFF files.
+        output_folder (str): The path to the folder to save the output JPG files.
+
+    Returns:
+        str: The path to the output folder containing the JPG files.
+    """
+
+    if output_folder is None:
+        output_folder = make_temp_dir()
+
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
+
+    geotiff_files = [
+        f for f in os.listdir(input_folder) if f.endswith(".tif") or f.endswith(".tiff")
+    ]
+
+    # Initialize tqdm progress bar
+    for filename in tqdm(geotiff_files, desc="Converting GeoTIFF to JPG"):
+        geotiff_path = os.path.join(input_folder, filename)
+        jpg_filename = os.path.splitext(filename)[0] + ".jpg"
+        output_path = os.path.join(output_folder, jpg_filename)
+        geotiff_to_jpg(geotiff_path, output_path)
+
+    return output_folder

samgeo/samgeo2.py

@@ -1071,15 +1071,14 @@ def set_video(
             step (int, optional): The step size. Defaults to 1.
             frame_rate (Optional[int], optional): The frame rate. Defaults to None.
         """
-        import tempfile
 
         if isinstance(video_path, str):
             if video_path.startswith("http"):
                 video_path = common.download_file(video_path)
             if os.path.isfile(video_path):
 
                 if output_dir is None:
-                    output_dir = tempfile.mkdtemp()
+                    output_dir = common.make_temp_dir()
                     if not os.path.exists(output_dir):
                         os.makedirs(output_dir)
                 print(f"Output directory: {output_dir}")
@@ -1088,6 +1087,14 @@ def set_video(
                 )
 
             elif os.path.isdir(video_path):
+                files = sorted(os.listdir(video_path))
+                if len(files) == 0:
+                    raise ValueError(f"No files found in {video_path}.")
+                elif files[0].endswith(".tif"):
+                    self._tif_source = os.path.join(video_path, files[0])
+                    self._tif_dir = video_path
+                    self._tif_names = files
+                    video_path = common.geotiff_to_jpg_batch(video_path)
                 output_dir = video_path
 
             if not os.path.exists(video_path):
@@ -1189,7 +1196,9 @@ def save_video_segments(self, output_dir: str, img_ext: str = "png") -> None:
         """
         from PIL import Image
 
-        def save_image_from_dict(data, output_path="output_image.png"):
+        def save_image_from_dict(
+            data, output_path="output_image.png", crs_source=None, **kwargs
+        ):
             # Find the shape of the first array in the dictionary (assuming all arrays have the same shape)
             array_shape = next(iter(data.values())).shape[1:]
 
@@ -1201,26 +1210,40 @@ def save_image_from_dict(data, output_path="output_image.png"):
                 # Assign the key value wherever the boolean array is True
                 output_array[array[0]] = key
 
-            # Convert the output array to a PIL image
-            image = Image.fromarray(output_array)
+            if crs_source is None:
+                # Convert the output array to a PIL image
+                image = Image.fromarray(output_array)
 
-            # Save the image
-            image.save(output_path)
-
-        if not os.path.exists(output_dir):
-            os.makedirs(output_dir)
+                # Save the image
+                image.save(output_path)
+            else:
+                output_path = output_path.replace(".png", ".tif")
+                common.array_to_image(output_array, output_path, crs_source, **kwargs)
 
         num_frames = len(self.video_segments)
         num_digits = len(str(num_frames))
 
+        if hasattr(self, "_tif_source") and self._tif_source.endswith(".tif"):
+            crs_source = self._tif_source
+            filenames = self._tif_names
+        else:
+            crs_source = None
+            filenames = None
+
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+
         # Initialize the tqdm progress bar
         for frame_idx, video_segment in tqdm(
             self.video_segments.items(), desc="Rendering frames", total=num_frames
         ):
-            output_path = os.path.join(
-                output_dir, f"{str(frame_idx).zfill(num_digits)}.{img_ext}"
-            )
-            save_image_from_dict(video_segment, output_path)
+            if filenames is None:
+                output_path = os.path.join(
+                    output_dir, f"{str(frame_idx).zfill(num_digits)}.{img_ext}"
+                )
+            else:
+                output_path = os.path.join(output_dir, filenames[frame_idx])
+            save_image_from_dict(video_segment, output_path, crs_source)
 
     def save_video_segments_blended(
         self,
@@ -1390,7 +1413,10 @@ def show_box(box, ax):
         prompts = self._convert_prompts(prompts)
         video_dir = self.video_path
         frame_names = self._frame_names
-        plt.figure(figsize=figsize)
+        fig = plt.figure(figsize=figsize)
+        fig.canvas.toolbar_visible = True
+        fig.canvas.header_visible = False
+        fig.canvas.footer_visible = True
         plt.title(f"frame {frame_idx}")
         plt.imshow(Image.open(os.path.join(video_dir, frame_names[frame_idx])))
 
@@ -1406,3 +1432,5 @@ def show_box(box, ax):
                     show_box(box, plt.gca())
                 if mask is not None:
                     show_mask(mask, plt.gca(), obj_id=obj_id)
+
+        plt.show()