Update app.py

app.py

@@ -597,16 +597,22 @@ with tab_ocr:
     st.header("Test OCR 🔍")
     all_files = get_gallery_files()
     if all_files:
+        # Filter for only PNG and PDF files
+        ocr_files = [f for f in all_files if f.endswith(('.png', '.pdf'))]
         if st.button("OCR All Assets 🚀"):
             full_text = "# OCR Results\n\n"
-            for file in all_files:
+            for file in ocr_files:
                 if file.endswith('.png'):
                     image = Image.open(file)
-                else:
-                    doc = fitz.open(file)
-                    pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
-                    image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-                    doc.close()
+                else:  # PDF
+                    try:
+                        doc = fitz.open(file)
+                        pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
+                        image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+                        doc.close()
+                    except Exception as e:
+                        st.error(f"Failed to process {file}: {str(e)}")
+                        continue
                 output_file = generate_filename(f"ocr_{os.path.basename(file)}", "txt")
                 result = asyncio.run(process_ocr(image, output_file))
                 full_text += f"## {os.path.basename(file)}\n\n{result}\n\n"
@@ -616,21 +622,28 @@ with tab_ocr:
                 f.write(full_text)
             st.success(f"Full OCR saved to {md_output_file}")
             st.markdown(get_download_link(md_output_file, "text/markdown", "Download Full OCR Markdown"), unsafe_allow_html=True)
-        selected_file = st.selectbox("Select Image or PDF", all_files, key="ocr_select")
+        selected_file = st.selectbox("Select Image or PDF", ocr_files, key="ocr_select")
         if selected_file:
             if selected_file.endswith('.png'):
                 image = Image.open(selected_file)
-            else:
-                doc = fitz.open(selected_file)
-                pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
-                image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-                doc.close()
-            st.image(image, caption="Input Image", use_container_width=True)
-            if st.button("Run OCR 🚀", key="ocr_run"):
-                output_file = generate_filename("ocr_output", "txt")
-                result = asyncio.run(process_ocr(image, output_file))
-                st.text_area("OCR Result", result, height=200)
-                st.session_state['history'].append(f"OCR Test: {selected_file} -> {output_file}")
+            else:  # PDF
+                try:
+                    doc = fitz.open(selected_file)
+                    pix = doc[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
+                    image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+                    doc.close()
+                except Exception as e:
+                    st.error(f"Cannot process {selected_file}: {str(e)}. Please select a PNG or PDF file.")
+                    image = None
+            if image:
+                st.image(image, caption="Input Image", use_container_width=True)
+                if st.button("Run OCR 🚀", key="ocr_run"):
+                    output_file = generate_filename("ocr_output", "txt")
+                    result = asyncio.run(process_ocr(image, output_file))
+                    st.text_area("OCR Result", result, height=200)
+                    st.session_state['history'].append(f"OCR Test: {selected_file} -> {output_file}")
+    else:
+        st.warning("No assets in gallery yet. Use Camera Snap or Download PDFs!")
 
 with tab_build:
     st.header("Build Titan 🌱")