Update app.py

app.py

@@ -4,6 +4,7 @@ import requests
 from PIL import Image
 from io import BytesIO
 import fitz  # PyMuPDF
+import os
 
 # Configuration - Replace with your API key
 GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
@@ -43,6 +44,22 @@ def encode_file(uploaded_file):
             img.save(img_byte_arr, format='JPEG')
             return base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
         
+        elif uploaded_file.type in ["image/webp", "image/avif", "image/jpeg", "image/png", "image/gif"]:
+            img = Image.open(BytesIO(file_bytes))
+            img_byte_arr = BytesIO()
+            img.save(img_byte_arr, format='JPEG')
+            return base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
+        
+        elif uploaded_file.type == "text/plain":
+            # Convert TXT to image
+            text = file_bytes.decode('utf-8')
+            img = Image.new('RGB', (800, 600), color = (73, 109, 137))
+            d = ImageDraw.Draw(img)
+            d.text((10,10), text, fill=(255,255,0))
+            img_byte_arr = BytesIO()
+            img.save(img_byte_arr, format='JPEG')
+            return base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
+        
         return base64.b64encode(file_bytes).decode('utf-8')
     except Exception as e:
         st.error(f"File processing error: {str(e)}")
@@ -94,8 +111,14 @@ def process_document():
                 pix = page.get_pixmap()
                 img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                 st.session_state.doc_preview = img
-            else:
+            elif uploaded_file.type in ["image/webp", "image/avif", "image/jpeg", "image/png", "image/gif"]:
                 st.session_state.doc_preview = Image.open(uploaded_file)
+            elif uploaded_file.type == "text/plain":
+                text = uploaded_file.getvalue().decode('utf-8')
+                img = Image.new('RGB', (800, 600), color = (73, 109, 137))
+                d = ImageDraw.Draw(img)
+                d.text((10,10), text, fill=(255,255,0))
+                st.session_state.doc_preview = img
 
             # Classify document
             classify_prompt = f"Classify this healthcare document into one of: {DOCUMENT_TYPES}. Respond only with the category name."
@@ -151,7 +174,7 @@ def main():
         # Upload button
         st.file_uploader(
             "Upload Document",
-            type=["pdf", "png", "jpg", "jpeg"],
+            type=["pdf", "webp", "avif", "jpg", "jpeg", "png", "gif", "txt"],
             key="uploaded_file",
             on_change=process_document
         )