Update app.py

app.py

@@ -4,6 +4,9 @@ import json
 import random
 from datetime import datetime
 import os
+from PIL import Image
+import io
+import numpy as np
 
 # Get access token from environment
 access_token = os.environ.get("HUGGINGFACE_TOKEN")
@@ -13,7 +16,31 @@ class DatasetViewer:
         self.dataset = None
         self.dataset_size = 0
         self.last_refresh_time = None
+        self.max_display_size = (800, 600)  # Maximum width and height for displayed images
         self.load_dataset()
+        
+    def resize_image(self, image):
+        """Resize image keeping aspect ratio with a maximum size constraint"""
+        if isinstance(image, np.ndarray):
+            # Convert numpy array to PIL Image
+            image = Image.fromarray(image)
+        elif isinstance(image, bytes):
+            # Convert bytes to PIL Image
+            image = Image.open(io.BytesIO(image))
+            
+        # Calculate scaling factor to fit within max dimensions
+        width_ratio = self.max_display_size[0] / image.width
+        height_ratio = self.max_display_size[1] / image.height
+        scale_factor = min(width_ratio, height_ratio)
+        
+        # Only resize if image is larger than max dimensions
+        if scale_factor < 1:
+            new_width = int(image.width * scale_factor)
+            new_height = int(image.height * scale_factor)
+            image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+            
+        # Convert back to numpy array for gradio
+        return np.array(image)
 
     def load_dataset(self):
         """Load the complete dataset into memory"""
@@ -53,8 +80,11 @@ class DatasetViewer:
             
             # Append the triple (post_info, source_image, edited_image)
             results.append(markdown_text)
-            results.append(sample["source_image"])
-            results.append(sample["edited_image"])
+            # Resize images before adding to results
+            source_image = self.resize_image(sample["source_image"])
+            edited_image = self.resize_image(sample["edited_image"])
+            results.append(source_image)
+            results.append(edited_image)
         
         return tuple(results)