Update open_tiff function for supporting images written by rasterio and tifffile

app.py

@@ -14,6 +14,7 @@ from mmseg.models import build_segmentor
 from mmseg.datasets.pipelines import Compose, LoadImageFromFile
 
 import rasterio
+import tifffile
 import torch
 
 from mmseg.apis import init_segmentor
@@ -45,13 +46,19 @@ def stretch_rgb(rgb):
     return img_rescale
 
 
-def open_tiff(fname):
-    
+def open_tiff(fname):  # it supports images written by rasterio and tifffile
     with rasterio.open(fname, "r") as src:
-        
-        data = src.read()
-        
-    return data
+        img_rasterio = src.read()
+
+    with tifffile.TiffFile(fname) as tif:
+        img_tifffile = tif.asarray()
+    
+    if img_rasterio.shape == img_tifffile.shape:
+        return img_rasterio
+    if img_rasterio.shape[0] == img_tifffile.shape[-1]:
+        return img_rasterio
+    return img_tifffile
+
 
 def write_tiff(img_wrt, filename, metadata):