Upload app.py

app.py

@@ -25,6 +25,8 @@ try:
     else:
         device = "cpu"
         dtype = torch.float32 # CPUs handle float32 best
+        # Prevent CPU thrashing on Spaces (usually 2 vCPUs)
+        torch.set_num_threads(4)
         print("Running on CPU mode")
     
     model_id = "lightonai/LightOnOCR-2-1B"
@@ -58,10 +60,10 @@ except Exception as e:
     processor = None
 
 # --- Helper Functions ---
-def resize_for_ocr(image, max_dim=896):
+def resize_for_ocr(image, max_dim=768):
     """
     Resize image to be faster.
-    Lowered max_dim from 1280 to 896 for CPU deployment to ensure reasonable speed.
+    Lowered max_dim from 1280->896->768 for CPU deployment to ensure reasonable speed.
     """
     if image is None: return None
     w, h = image.size
@@ -229,6 +231,15 @@ def stream_ocr(image):
     except Exception as e:
         yield f"Error during processing: {str(e)}", None
 
+# --- Prepare Examples ---
+example_images = []
+data_dir = os.path.join(os.path.dirname(__file__), 'data')
+if os.path.exists(data_dir):
+    valid_exts = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
+    # Gradio Examples expects a list of lists: [[path1], [path2], ...]
+    example_images = [[os.path.join(data_dir, f)] for f in os.listdir(data_dir) 
+                      if os.path.splitext(f)[1].lower() in valid_exts]
+
 # --- Aesthetic Custom CSS ---
 custom_css = """
 /* Dark Purple Gradient Background */
@@ -343,6 +354,17 @@ with gr.Blocks(title="Ultra OCR") as demo:
                         gr.Markdown("### Download Results")
                         output_file = gr.File(label="Download Word (.docx)", type="filepath")
 
+        # Example Gallery
+        if example_images:
+            gr.HTML("<hr>")
+            gr.Markdown("### 📂 Sample Documents")
+            gr.Examples(
+                examples=example_images,
+                inputs=input_img,
+                label="Click a sample to test",
+                examples_per_page=5
+            )
+
     # Interactions
     run_btn.click(
         fn=stream_ocr,