Update app.py

app.py

@@ -1,3 +1,84 @@
+'''
+from datasets import load_dataset
+from gradio_client import Client, handle_file
+import os
+from PIL import Image
+import tempfile
+from tqdm import tqdm
+
+# Load the dataset
+ds = load_dataset("svjack/Aesthetics_X_Phone_4K_Images_Rec_Captioned")
+
+# Initialize Gradio client
+client = Client("http://localhost:7860")
+
+# Create output directory if it doesn't exist
+output_dir = "processed_output"
+os.makedirs(output_dir, exist_ok=True)
+
+# Iterate through all items in the training set
+for idx, item in tqdm(enumerate(ds["train"])):
+    try:
+        image = item["image"]
+        joy_caption = item["joy-caption"]
+        
+        # Create a temporary file for the input image
+        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
+            temp_image_path = temp_file.name
+            image.save(temp_image_path)
+        
+        # Process the image through the API
+        result = client.predict(
+            image=handle_file(temp_image_path),
+            width=1280,
+            height=720,
+            overlap_percentage=10,
+            num_inference_steps=8,
+            resize_option="Full",
+            custom_resize_percentage=50,
+            prompt_input="",
+            alignment="Middle",
+            overlap_left=True,
+            overlap_right=True,
+            overlap_top=True,
+            overlap_bottom=True,
+            api_name="/infer"
+        )
+        
+        # Get the processed image path from the result
+        processed_image_path = result[1]
+        
+        # Define output paths
+        base_filename = f"processed_{idx}"
+        output_image_path = os.path.join(output_dir, f"{base_filename}.png")
+        output_text_path = os.path.join(output_dir, f"{base_filename}.txt")
+        
+        # Ensure the output is saved as PNG
+        if processed_image_path.lower().endswith('.png'):
+            # If already PNG, just copy
+            with Image.open(processed_image_path) as img:
+                img.save(output_image_path, 'PNG')
+        else:
+            # If not PNG, open and convert to PNG
+            with Image.open(processed_image_path) as img:
+                img.save(output_image_path, 'PNG')
+        
+        # Save the joy-caption as a text file
+        with open(output_text_path, 'w', encoding='utf-8') as f:
+            f.write(joy_caption)
+            
+        print(f"Processed item {idx}: Image saved to {output_image_path}, caption to {output_text_path}")
+        
+    except Exception as e:
+        print(f"Error processing item {idx}: {str(e)}")
+    finally:
+        # Clean up temporary files
+        if 'temp_image_path' in locals() and os.path.exists(temp_image_path):
+            os.remove(temp_image_path)
+
+print("Processing complete!")
+'''
+
 import gradio as gr
 import spaces
 import torch