Update

.ipynb_checkpoints/app-checkpoint.py

@@ -3,10 +3,18 @@ import gradio as gr
 import re
 import json
 import os
-
+from pdf2image import convert_from_bytes
 from transformers import DonutProcessor, VisionEncoderDecoderModel
 from datasets import load_dataset 
 import torch
+import imghdr
+
+def check_image(image):
+    if imghdr.what(image) == 'jpeg' or imghdr.what(image) == 'png':
+        return image
+    else:
+        images = convert_from_bytes(image.read(), fmt="jpeg", size=(960,1280))
+    return images[0]
 
 def get_attributes(input_img):
     access_token = str(os.environ.get('key'))
@@ -17,8 +25,9 @@ def get_attributes(input_img):
 
     model.eval()
     model.to(device)
-        
-    pixel_values = processor(input_img, return_tensors="pt").pixel_values
+    
+    image = check_image(input_img)
+    pixel_values = processor(image, return_tensors="pt").pixel_values
     pixel_values = pixel_values.to(device)
     print(pixel_values.size())
     # prepare decoder inputs

.ipynb_checkpoints/requirements-checkpoint.txt

@@ -1,4 +1,6 @@
 torch
 transformers
 numpy
-datasets
+datasets
+pdf2image
+imghdr

app.py

@@ -3,10 +3,18 @@ import gradio as gr
 import re
 import json
 import os
-
+from pdf2image import convert_from_bytes
 from transformers import DonutProcessor, VisionEncoderDecoderModel
 from datasets import load_dataset 
 import torch
+import imghdr
+
+def check_image(image):
+    if imghdr.what(image) == 'jpeg' or imghdr.what(image) == 'png':
+        return image
+    else:
+        images = convert_from_bytes(image.read(), fmt="jpeg", size=(960,1280))
+    return images[0]
 
 def get_attributes(input_img):
     access_token = str(os.environ.get('key'))
@@ -17,8 +25,9 @@ def get_attributes(input_img):
 
     model.eval()
     model.to(device)
-        
-    pixel_values = processor(input_img, return_tensors="pt").pixel_values
+    
+    image = check_image(input_img)
+    pixel_values = processor(image, return_tensors="pt").pixel_values
     pixel_values = pixel_values.to(device)
     print(pixel_values.size())
     # prepare decoder inputs

requirements.txt

@@ -1,4 +1,6 @@
 torch
 transformers
 numpy
-datasets
+datasets
+pdf2image
+imghdr