Update app.py

app.py

@@ -2758,41 +2758,55 @@ def save_pdf_to_generated_dir(pdf_stream: io.BytesIO, project_id: str) -> str:
         logger.error(f"Failed to save PDF to generated dir: {e}", exc_info=True)
         return None
 
-def pdf_to_images_with_size_check(pdf_path, output_dir, size_limit_mb=4):
-    os.makedirs(output_dir, exist_ok=True)
-
-    # Convert PDF to images
-    images = convert_from_path(pdf_path, dpi=300)  # 300 DPI keeps quality
-    saved_files = []
-
-    for i, img in enumerate(images, start=1):
-        output_path = os.path.join(output_dir, f"page_{i}.jpg")
-
-        # Save to memory first to check size
-        img_bytes = io.BytesIO()
-        img.save(img_bytes, format="JPEG", quality=95)  # near-lossless
-        size_mb = len(img_bytes.getvalue()) / (1024 * 1024)
-
-        if size_mb > size_limit_mb:
-            print(f"Page {i}: {size_mb:.2f} MB → compressing...")
-            # Compress until under size limit
-            quality = 95
-            while size_mb > size_limit_mb and quality > 70:  # don't go below 70
-                img_bytes = io.BytesIO()
-                img.save(img_bytes, format="JPEG", quality=quality)
-                size_mb = len(img_bytes.getvalue()) / (1024 * 1024)
-                quality -= 5
-        else:
-            print(f"Page {i}: {size_mb:.2f} MB → no compression needed.")
-
-        # Write final image to disk
-        with open(output_path, "wb") as f:
-            f.write(img_bytes.getvalue())
-
-        saved_files.append(output_path)
-
-    return saved_files
+# def pdf_to_images_with_size_check(pdf_path, output_dir, size_limit_mb=4):
+#     os.makedirs(output_dir, exist_ok=True)
+
+#     # Convert PDF to images
+#     images = convert_from_path(pdf_path, dpi=300)  # 300 DPI keeps quality
+#     saved_files = []
+
+#     for i, img in enumerate(images, start=1):
+#         output_path = os.path.join(output_dir, f"page_{i}.jpg")
+
+#         # Save to memory first to check size
+#         img_bytes = io.BytesIO()
+#         img.save(img_bytes, format="JPEG", quality=95)  # near-lossless
+#         size_mb = len(img_bytes.getvalue()) / (1024 * 1024)
+
+#         if size_mb > size_limit_mb:
+#             print(f"Page {i}: {size_mb:.2f} MB → compressing...")
+#             # Compress until under size limit
+#             quality = 95
+#             while size_mb > size_limit_mb and quality > 70:  # don't go below 70
+#                 img_bytes = io.BytesIO()
+#                 img.save(img_bytes, format="JPEG", quality=quality)
+#                 size_mb = len(img_bytes.getvalue()) / (1024 * 1024)
+#                 quality -= 5
+#         else:
+#             print(f"Page {i}: {size_mb:.2f} MB → no compression needed.")
+
+#         # Write final image to disk
+#         with open(output_path, "wb") as f:
+#             f.write(img_bytes.getvalue())
+
+#         saved_files.append(output_path)
+
+#     return saved_files
+def compress_image_if_needed(image, max_size_mb=4, quality=85):
+    """
+    Compress the given PIL Image if its size is greater than max_size_mb.
+    Returns the (possibly compressed) image object.
+    """
+    temp_path = "/tmp/temp_compression_check.jpg"
+    image.save(temp_path, format="JPEG", quality=95)  # save original temporarily
+    size_mb = os.path.getsize(temp_path) / (1024 * 1024)
 
+    if size_mb > max_size_mb:
+        # Compress by reducing quality
+        image.save(temp_path, format="JPEG", quality=quality, optimize=True)
+        print(f"Image compressed from {size_mb:.2f} MB to {os.path.getsize(temp_path)/(1024*1024):.2f} MB")
+        return Image.open(temp_path)
+    return image
 
 @app.route('/')
 def index():
@@ -2871,9 +2885,9 @@ def process_pdf():
         # }
 
         # Save uploaded file to disk
-        pdf_path = os.path.join("/tmp", secure_filename(file.filename))
-        file.save(pdf_path)        
-        compressed_pages = pdf_to_images_with_size_check(pdf_path, "/tmp/compressed_pages", size_limit_mb=4)
+        # pdf_path = os.path.join("/tmp", secure_filename(file.filename))
+        # file.save(pdf_path)        
+        # compressed_pages = pdf_to_images_with_size_check(pdf_path, "/tmp/compressed_pages", size_limit_mb=4)
         
         # {
         # Extract & process
@@ -2921,6 +2935,12 @@ def process_pdf():
         else:
             images = convert_from_path(pdf_stream, dpi=300)
 
+        # Compress images if needed
+        compressed_images = []
+        for img in images:
+            compressed_images.append(compress_image_if_needed(img, max_size_mb=4, quality=85))
+        images = compressed_images
+        
         #updating logic here [Dev Patel]
         initial_state_dict = {
             "project_json": project_skeleton,