Update main.py

main.py

@@ -76,6 +76,7 @@ import gc
 import logging
 from typing import List
 import time
+import numpy as np
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
@@ -219,9 +220,17 @@ async def batch_extract_text(batch_size:int, background_tasks: BackgroundTasks,
         logger.error(f"Error in batch processing: {str(e)}")
         return {"error": str(e)}
 
+
 # Health check endpoint
 @app.get("/health")
 async def health_check():
+    # Generate a random image (20x40 pixels) with random RGB values
+    random_data = np.random.randint(0, 256, (20, 40, 3), dtype=np.uint8)
+    
+    # Create an image from the random data
+    image = Image.fromarray(random_data)
+    predict(image)
+    clean_memory()
     return {"status": "healthy"}
 
 # if __name__ == "__main__":