Create app.py

app.py

@@ -0,0 +1,119 @@
+import os
+from PIL import Image
+import time
+from gradio_client import Client, handle_file
+import zipfile
+import gradio as gr
+
+# Configuration
+INPUT_DIR = "input_images"  # Folder with original images
+OUTPUT_DIR = "output_images"  # Base output folder
+DATASET_DIR = os.path.join(OUTPUT_DIR, "dataset")  # Subfolder for organized dataset
+ZIP_FILE = os.path.join(OUTPUT_DIR, "dataset.zip")  # Path for the output ZIP file
+TARGET_SIZE = (512, 512)  # Target size for Stable Diffusion (SD 1.5)
+HUGGINGFACE_SPACE_URL = "bdsqlsz/Florence-2-SD3-Captioner"
+
+# Ensure output and dataset directories exist
+os.makedirs(DATASET_DIR, exist_ok=True)
+
+def resize_and_crop_image(input_path, output_path, target_size):
+    """Resize and crop an image to the target size while preserving aspect ratio."""
+    try:
+        img = Image.open(input_path).convert("RGB")
+        width, height = img.size
+        target_width, target_height = target_size
+        img_ratio = width / height
+        target_ratio = target_width / target_height
+
+        if img_ratio > target_ratio:
+            new_height = target_height
+            new_width = int(new_height * img_ratio)
+        else:
+            new_width = target_width
+            new_height = int(new_width / img_ratio)
+
+        img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+        left = (new_width - target_width) // 2
+        top = (new_height - target_height) // 2
+        right = left + target_width
+        bottom = top + target_height
+        img = img.crop((left, top, right, bottom))
+        img.save(output_path, "JPEG", quality=95)
+        return True
+    except Exception as e:
+        print(f"Error processing {input_path}: {e}")
+        return False
+
+def get_caption_from_florence(image_path):
+    """Call the Florence-2-SD3-Captioner /process_image endpoint via Gradio API."""
+    try:
+        client = Client(HUGGINGFACE_SPACE_URL)
+        result = client.predict(
+            image=handle_file(image_path),
+            api_name="/process_image"
+        )
+        return result if isinstance(result, str) else "No caption returned"
+    except Exception as e:
+        print(f"Error captioning {image_path}: {e}")
+        return "Captioning failed"
+
+def create_zip_file():
+    """Create a ZIP file of the dataset folder."""
+    with zipfile.ZipFile(ZIP_FILE, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        for root, _, files in os.walk(DATASET_DIR):
+            for file in files:
+                file_path = os.path.join(root, file)
+                arcname = os.path.relpath(file_path, OUTPUT_DIR)
+                zipf.write(file_path, arcname)
+    return ZIP_FILE
+
+def process_images():
+    """Process all images and return status."""
+    if not os.path.exists(INPUT_DIR):
+        return f"Input directory '{INPUT_DIR}' not found."
+    
+    image_files = [f for f in os.listdir(INPUT_DIR) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
+    if not image_files:
+        return f"No images found in '{INPUT_DIR}'."
+
+    for idx, filename in enumerate(image_files, 1):
+        input_path = os.path.join(INPUT_DIR, filename)
+        base_name = f"img{idx}"
+        output_image_path = os.path.join(DATASET_DIR, f"{base_name}.jpg")
+        caption_file_path = os.path.join(DATASET_DIR, f"{base_name}.txt")
+
+        print(f"Processing {idx}/{len(image_files)}: {filename}")
+
+        if resize_and_crop_image(input_path, output_image_path, TARGET_SIZE):
+            caption = get_caption_from_florence(output_image_path)
+            print(f"Caption: {caption}")
+            with open(caption_file_path, "w", encoding="utf-8") as f:
+                f.write(caption)
+        else:
+            print(f"Skipping captioning for {filename} due to processing error.")
+        
+        time.sleep(1)  # Avoid overwhelming the API
+    
+    # Create ZIP file after processing
+    zip_path = create_zip_file()
+    return f"Processing complete! ZIP file created at {zip_path}"
+
+def launch_interface():
+    """Launch Gradio interface with a download button after processing."""
+    status = process_images()
+    
+    with gr.Blocks(title="Image Processing and Download") as demo:
+        gr.Markdown("### Image Processing Status")
+        status_text = gr.Textbox(value=status, label="Status", interactive=False)
+        
+        if "Processing complete" in status:
+            gr.Markdown("### Download Your Dataset")
+            download_button = gr.File(label="Download ZIP", value=ZIP_FILE)
+        else:
+            gr.Markdown("No ZIP file available due to processing errors.")
+        
+        demo.launch()
+
+if __name__ == "__main__":
+    print("Starting image processing and captioning...")
+    launch_interface()