Fix app.py

app.py

@@ -40,16 +40,25 @@ structure_class_thresholds = {
 }
 
 
+def PIL_to_cv(pil_img):
+    return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR) 
+
+
+def cv_to_PIL(cv_img):
+    return PIL.Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
+
+
 def table_structure(filename):
-    image = cv2.imread(filename)
+    pil_img = PIL.Image.open(filename)
+    image = PIL_to_cv(pil_img)
     pred = structure_model(image, size=imgsz)
-    pred = pred.xywhn[0]    
+    pred = pred.xywhn[0]
     result = pred.cpu().numpy()
     return result
 
 
 def ocr(filename):
-    doc = DocumentFile.from_images(filename)
+    doc = DocumentFile.from_images(filename.read())
     result = ocr_predictor(doc).export()
     result = result['pages'][0]
     H, W = result['dimensions']
@@ -67,7 +76,9 @@ def ocr(filename):
 
 
 def convert_stucture(page_tokens, filename, structure_result):
-    image = cv2.imread(filename)
+    pil_img = PIL.Image.open(filename)
+    image = PIL_to_cv(pil_img)
+
     width = image.shape[1]
     height = image.shape[0]
     # print(width, height)
@@ -119,7 +130,8 @@ def convert_stucture(page_tokens, filename, structure_result):
 
 
 def visualize_cells(filename, cells, ax):    
-    image = cv2.imread(filename)
+    pil_img = PIL.Image.open(filename)
+    image = PIL_to_cv(pil_img)
     for i, cell in enumerate(cells):
         bbox = cell['bbox']
         x1 = int(bbox[0])
@@ -127,7 +139,7 @@ def visualize_cells(filename, cells, ax):
         x2 = int(bbox[2])
         y2 = int(bbox[3])
         cv2.rectangle(image, (x1, y1), (x2, y2), color=(0, 255, 0))
-    ax.image(image)
+    ax.image(cv_to_PIL(image))
 
 
 def pytess(cell_pil_img):
@@ -234,7 +246,7 @@ def main():
         else:
             print(filename)
 
-            cols[0].image(cv2.imread(filename))
+            cols[0].image(filename)
 
             ocr_res = ocr(filename)
             structure_result = table_structure(filename)