Update app.py

app.py

@@ -1,6 +1,8 @@
 import streamlit as st
-import imageio.v2
 from transformers import GitProcessor, GitForCausalLM, set_seed
+import tifffile as tiff
+import numpy as np
+from PIL import Image
 
 set_seed(226)
 model = GitForCausalLM.from_pretrained("microsoft/git-base").to("cuda").eval()
@@ -10,9 +12,12 @@ st.title("What is hidden ?")
 file_name = st.file_uploader("Upload your image")
 if file_name is not None:
     col1, col2 = st.columns(2)
-    image = imageio.v2.imread(file_name)
-    inputs = processor(images=image, return_tensors="pt", do_rescale=file_name.endswith(".tiff"), do_resize=False).to("cuda")
+    isTIFF = file_name.name.endswith(".tiff")
+    image = tiff.imread(file_name) if isTIFF else Image.open(file_name)
+    image = np.array(image)
+    inputs = processor(images=image, return_tensors="pt", do_rescale=not isTIFF, do_resize=False).to("cuda")
     generated_ids = model.generate(**inputs)
     generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    col1.image(image, use_column_width=True)
-    col2.header(generated_caption)
+    print(generated_caption)
+    col1.image(image, use_container_width=True)
+    col2.header(generated_caption)