Update main.py

main.py

@@ -8,6 +8,29 @@ import string
 from mimetypes import guess_type
 import requests
 import time
+from PIL import Image
+
+def compress_image(input_stream: io.BytesIO, max_size_mb=4) -> io.BytesIO:
+    with Image.open(input_stream) as img:
+        img = Image.new(img.mode, img.size)
+        img.paste(img)
+        quality = 95
+        step = 5
+        output_stream = io.BytesIO()
+        while True:
+            output_stream.seek(0)
+            img.save(output_stream, format="JPEG", quality=quality)
+            output_stream_size_mb = len(output_stream.getvalue()) / (1024 * 1024)
+
+            if output_stream_size_mb <= max_size_mb:
+                output_stream.seek(0)
+                print(f"Image compressed to {output_stream_size_mb:.2f} MB")
+                return output_stream
+            quality -= step
+            if quality < step:
+                print("Cannot compress the image to the desired size.")
+                output_stream.seek(0)
+                return output_stream
 
 def read_image_from_stream(image_data, subscription_key, endpoint):
 
@@ -109,6 +132,7 @@ def valuate_qna(p: str, r: str) -> str:
     return hug.output.get("data", [None])[0]
 
 def ocr(image: io.BytesIO) -> str:
+    image = compress_image(image) # Compress the image
     return process_image(image, os.getenv('AZKEY'), os.getenv('AZURL'))
 
 app = Flask(__name__)