reduce image size if too large

app.py

@@ -13,6 +13,56 @@ logging.basicConfig(level=logging.ERROR)
 logger = logging.getLogger(__name__)
 
 
+def resize_image_if_needed(pil_image, max_size_mb=1, max_edge_length=1024):
+    """
+    Detect the size of a PIL image, and if it exceeds 1MB or its long edge is larger than 1024 pixels,
+    reduce its size to a smaller size.
+    
+    Args:
+        pil_image (PIL.Image.Image): The input PIL image.
+        max_size_mb (int): The maximum allowed size in megabytes.
+        max_edge_length (int): The maximum allowed length of the long edge in pixels.
+    
+    Returns:
+        PIL.Image.Image: The resized PIL image.
+    """
+    # Convert image to bytes and check its size
+    img_byte_arr = io.BytesIO()
+    pil_image.save(img_byte_arr, format='JPEG')
+    img_size_mb = len(img_byte_arr.getvalue()) / (1024 * 1024)
+    print(f"Image size: {img_size_mb} MB")
+    
+    # Check if the image size exceeds the maximum allowed size
+    if img_size_mb > max_size_mb or max(pil_image.size) > max_edge_length:
+        # Calculate the new size while maintaining the aspect ratio
+        aspect_ratio = pil_image.width / pil_image.height
+        if pil_image.width > pil_image.height:
+            new_width = min(max_edge_length, pil_image.width)
+            new_height = int(new_width / aspect_ratio)
+        else:
+            new_height = min(max_edge_length, pil_image.height)
+            new_width = int(new_height * aspect_ratio)
+        
+        # Resize the image
+        pil_image = pil_image.resize((new_width, new_height), Image.LANCZOS)
+        
+        # Convert the resized image to bytes and check its size again
+        img_byte_arr = io.BytesIO()
+        pil_image.save(img_byte_arr, format='JPEG')
+        img_size_mb = len(img_byte_arr.getvalue()) / (1024 * 1024)
+        
+        # If the resized image still exceeds the maximum allowed size, reduce the quality
+        if img_size_mb > max_size_mb:
+            quality = 95
+            while img_size_mb > max_size_mb and quality > 10:
+                img_byte_arr = io.BytesIO()
+                pil_image.save(img_byte_arr, format='JPEG', quality=quality)
+                img_size_mb = len(img_byte_arr.getvalue()) / (1024 * 1024)
+                quality -= 5
+                
+    return pil_image
+
+
 def clone_repo():
     # 从环境变量中获取 GitHub Token
     github_token = os.getenv('GH_TOKEN')
@@ -103,7 +153,8 @@ if clone_repo():
         pil_image = get_page_image(in_file, page_number)
     else:
         pil_image = Image.open(in_file).convert("RGB")
-
+    pil_image = resize_image_if_needed(pil_image)
+    
     text_rec = st.sidebar.button("認識開始")
 
     if pil_image is None: