import gradio as gr import os from PIL import Image, ImageOps, UnidentifiedImageError import tempfile import shutil from pathlib import Path import zipfile import json import math class ImagePrepApp: def __init__(self): self.images = [] self.current_index = 0 self.temp_dir = tempfile.mkdtemp() self.output_dir = os.path.join(self.temp_dir, "crops") self.thumbnails_dir = os.path.join(self.temp_dir, "thumbnails") self.display_dir = os.path.join(self.temp_dir, "display") self.selected_for_deletion = set() # Track selected files for deletion os.makedirs(self.output_dir, exist_ok=True) os.makedirs(self.thumbnails_dir, exist_ok=True) os.makedirs(self.display_dir, exist_ok=True) # Preset crop dimensions self.crop_presets = { "512x512": (512, 512), "768x768": (768, 768), "1024x1024": (1024, 1024), "2048x2048": (2048, 2048), "512x768": (512, 768), "768x512": (768, 512), "Custom": (0, 0) } # Display size options self.display_sizes = { "Small (600x480)": (600, 480), "Medium (800x600)": (800, 600), "Large (1000x750)": (1000, 750), "X-Large (1200x900)": (1200, 900), "Original Size": (0, 0) } # Current zoom and crop settings self.current_zoom = 1.0 self.current_crop_width = 512 self.current_crop_height = 512 # Utilities processing directory self.utilities_dir = os.path.join(self.temp_dir, "utilities") self.processed_dir = os.path.join(self.utilities_dir, "processed") self.corrupted_dir = os.path.join(self.utilities_dir, "corrupted") os.makedirs(self.utilities_dir, exist_ok=True) os.makedirs(self.processed_dir, exist_ok=True) os.makedirs(self.corrupted_dir, exist_ok=True) def calculate_gallery_height(self, num_images, columns=8): """Calculate optimal gallery height based on number of images""" # For galleries to scroll properly in Gradio, we need consistent heights # Return a reasonable fixed height that works for most cases if num_images <= 8: # 1 row return 150 elif num_images <= 16: # 2 rows return 250 elif num_images <= 24: # 3 rows return 350 else: # 4+ rows - use scrolling return 400 def calculate_output_gallery_height(self, num_images, columns=6): """Calculate optimal output gallery height based on number of images""" # For galleries to scroll properly in Gradio, we need consistent heights if num_images <= 6: # 1 row return 150 elif num_images <= 12: # 2 rows return 250 elif num_images <= 18: # 3 rows return 350 else: # 4+ rows - use scrolling return 400 def create_thumbnail(self, image_path, size=(150, 150)): """Create thumbnail for gallery display""" try: with Image.open(image_path) as img: img.thumbnail(size, Image.Resampling.LANCZOS) thumb_filename = f"thumb_{os.path.basename(image_path)}" thumb_path = os.path.join(self.thumbnails_dir, thumb_filename) img.save(thumb_path, "JPEG", quality=85) return thumb_path except Exception as e: print(f"Error creating thumbnail: {e}") return image_path def create_display_image(self, image_path, display_size_name="Medium (800x600)"): """Create display-sized image based on selected display size""" try: with Image.open(image_path) as img: if display_size_name == "Original Size": # Return original image display_filename = f"display_{os.path.basename(image_path)}" display_path = os.path.join(self.display_dir, display_filename) img.save(display_path, "JPEG", quality=95) return display_path, img.size else: # Resize to fit within specified size max_size = self.display_sizes[display_size_name] img.thumbnail(max_size, Image.Resampling.LANCZOS) display_filename = f"display_{os.path.basename(image_path)}" display_path = os.path.join(self.display_dir, display_filename) img.save(display_path, "JPEG", quality=90) return display_path, img.size except Exception as e: print(f"Error creating display image: {e}") return image_path, (0, 0) def load_images_from_folder(self, files): """Load images from uploaded files""" if not files: return "No files uploaded", [], gr.update() self.images = [] supported_formats = ('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp') for file in files: if file.name.lower().endswith(supported_formats): self.images.append(file.name) if not self.images: return "No supported image files found", [], gr.update() self.current_index = 0 # Create thumbnails for gallery thumbnail_paths = [] for img_path in self.images: thumb_path = self.create_thumbnail(img_path) thumbnail_paths.append(thumb_path) return ( f"āœ… Successfully loaded {len(self.images)} images", thumbnail_paths, gr.update(visible=True) # Show cropping tab ) def select_from_gallery(self, evt: gr.SelectData, display_size_name): """Select image from thumbnail gallery""" if evt.index < len(self.images): self.current_index = evt.index self.current_zoom = 1.0 # Reset zoom # Create display-sized version for cropping interface original_path = self.images[self.current_index] display_path, display_size = self.create_display_image(original_path, display_size_name) # Load display image display_image = Image.open(display_path) return ( display_image, f"Image {self.current_index + 1} of {len(self.images)} - Display: {display_size[0]}x{display_size[1]}", 1.0 # Reset zoom slider ) return None, "0/0", 1.0 def update_display_size(self, display_size_name): """Update display size when dropdown changes""" if not self.images: return None, "No images loaded" # Recreate display image with new size original_path = self.images[self.current_index] display_path, display_size = self.create_display_image(original_path, display_size_name) display_image = Image.open(display_path) return ( display_image, f"Image {self.current_index + 1} of {len(self.images)} - Display: {display_size[0]}x{display_size[1]}" ) def update_crop_dimensions(self, preset_choice, custom_width, custom_height): """Update crop dimensions based on preset selection""" if preset_choice == "Custom": self.current_crop_width = int(custom_width) if custom_width > 0 else 512 self.current_crop_height = int(custom_height) if custom_height > 0 else 512 return gr.update(visible=True), custom_width, custom_height else: width, height = self.crop_presets[preset_choice] self.current_crop_width = width self.current_crop_height = height return gr.update(visible=False), width, height def update_zoom(self, zoom_value): """Update zoom level for cropping""" # Ensure zoom_value is a float, not a string try: zoom_val = float(zoom_value) self.current_zoom = zoom_val return f"Zoom: {zoom_val:.1f}x" except (ValueError, TypeError): # If conversion fails, return default self.current_zoom = 1.0 return "Zoom: 1.0x" def navigate_image(self, direction, display_size_name="Medium (800x600)"): """Navigate to next or previous image""" if not self.images: return None, "No images loaded", 1.0 if direction == "next": self.current_index = (self.current_index + 1) % len(self.images) elif direction == "prev": self.current_index = (self.current_index - 1) % len(self.images) self.current_zoom = 1.0 # Reset zoom # Create display-sized version for cropping interface original_path = self.images[self.current_index] display_path, display_size = self.create_display_image(original_path, display_size_name) # Load display image display_image = Image.open(display_path) return ( display_image, f"Image {self.current_index + 1} of {len(self.images)} - Display: {display_size[0]}x{display_size[1]}", 1.0 # Reset zoom slider ) def toggle_gallery_drawer(self, current_visibility): """Toggle the visibility of the gallery drawer""" return not current_visibility def process_crop_click(self, image, crop_preset, custom_width, custom_height, zoom_value, display_size_name, evt: gr.SelectData): """Process crop when user clicks on image with zoom consideration""" if image is None or not self.images: return None, "No image loaded" try: # Get crop dimensions if crop_preset == "Custom": base_crop_width = int(custom_width) if custom_width > 0 else 100 base_crop_height = int(custom_height) if custom_height > 0 else 100 else: base_crop_width, base_crop_height = self.crop_presets[crop_preset] # Apply zoom to crop dimensions (CORRECTED: higher zoom = smaller crop area) effective_crop_width = int(base_crop_width / zoom_value) effective_crop_height = int(base_crop_height / zoom_value) # Get click coordinates from display image click_x = evt.index[0] if evt.index else 0 click_y = evt.index[1] if evt.index else 0 # Load original full-resolution image original_image = Image.open(self.images[self.current_index]) orig_width, orig_height = original_image.size # Get display image dimensions display_width, display_height = image.size # Calculate scale factors scale_x = orig_width / display_width scale_y = orig_height / display_height # Convert click coordinates to original image coordinates orig_click_x = int(click_x * scale_x) orig_click_y = int(click_y * scale_y) # Center crop box on click point in original coordinates crop_x = max(0, min(orig_click_x - effective_crop_width // 2, orig_width - effective_crop_width)) crop_y = max(0, min(orig_click_y - effective_crop_height // 2, orig_height - effective_crop_height)) # Ensure crop dimensions fit within original image actual_crop_width = min(effective_crop_width, orig_width - crop_x) actual_crop_height = min(effective_crop_height, orig_height - crop_y) # Crop from original full-resolution image cropped = original_image.crop((crop_x, crop_y, crop_x + actual_crop_width, crop_y + actual_crop_height)) # Resize to target dimensions cropped = cropped.resize((base_crop_width, base_crop_height), Image.Resampling.LANCZOS) zoom_info = f" (Zoom: {zoom_value:.1f}x)" if zoom_value != 1.0 else "" return cropped, f"Cropped: {base_crop_width}x{base_crop_height} from ({crop_x}, {crop_y}){zoom_info}" except Exception as e: return None, f"Error cropping image: {str(e)}" def save_crop(self, cropped_image): """Save cropped image to output directory""" if cropped_image is None: return "No cropped image to save", gr.update() try: # Generate filename base_name = os.path.splitext(os.path.basename(self.images[self.current_index]))[0] crop_count = len([f for f in os.listdir(self.output_dir) if f.startswith(base_name)]) + 1 output_filename = f"{base_name}_crop_{crop_count}.png" output_path = os.path.join(self.output_dir, output_filename) # Save the image cropped_image.save(output_path, "PNG") # Check if this is the first crop saved - if so, make download tab visible total_crops = len([f for f in os.listdir(self.output_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]) tab_update = gr.update(visible=True) if total_crops == 1 else gr.update() return f"āœ… Saved as {output_filename}", tab_update except Exception as e: return f"āŒ Error saving image: {str(e)}", gr.update() def get_output_gallery(self): """Get list of output images for gallery""" if not os.path.exists(self.output_dir): return [] output_files = [] for filename in sorted(os.listdir(self.output_dir)): if filename.lower().endswith(('.png', '.jpg', '.jpeg')): file_path = os.path.join(self.output_dir, filename) output_files.append(file_path) return output_files def create_clean_thumbnail(self, image_path, max_thumb_size=150): """Create clean thumbnail without text overlay for output gallery""" try: with Image.open(image_path) as img: # Get original dimensions for caption width, height = img.size # Calculate thumbnail size while preserving aspect ratio aspect_ratio = width / height if aspect_ratio > 1: # Wider than tall thumb_width = max_thumb_size thumb_height = int(max_thumb_size / aspect_ratio) else: # Taller than wide or square thumb_height = max_thumb_size thumb_width = int(max_thumb_size * aspect_ratio) # Resize image to calculated thumbnail size img_resized = img.resize((thumb_width, thumb_height), Image.Resampling.LANCZOS) # Save clean thumbnail base_filename = os.path.splitext(os.path.basename(image_path))[0] clean_filename = f"clean_{base_filename}.jpg" clean_path = os.path.join(self.thumbnails_dir, clean_filename) img_resized.save(clean_path, "JPEG", quality=85) return clean_path, f"{width}Ɨ{height}" except Exception as e: print(f"Error creating clean thumbnail: {e}") return image_path, "Error" def toggle_file_selection(self, evt: gr.SelectData): """Toggle selection of a file for deletion when left-clicked""" try: # Get original file paths (not the display tuples) original_files = self.get_output_gallery() if evt.index < len(original_files): file_path = original_files[evt.index] filename = os.path.basename(file_path) if file_path in self.selected_for_deletion: self.selected_for_deletion.remove(file_path) status = f"šŸ”µ Deselected: {filename}" else: self.selected_for_deletion.add(file_path) status = f"šŸ”“ Selected for deletion: {filename}" # Return updated gallery with visual indicators updated_gallery = self.get_output_gallery_with_selection_visual() selected_names = [os.path.basename(f) for f in self.selected_for_deletion] selection_text = f"Selected ({len(selected_names)}): {', '.join(selected_names) if selected_names else 'None'}" return list(updated_gallery), selection_text, status # If click failed, return current state current_gallery = self.get_output_gallery_with_selection_visual() return list(current_gallery), "No file selected", "Click failed" except Exception as e: current_gallery = self.get_output_gallery_with_selection_visual() return list(current_gallery), f"Error: {str(e)}", f"Error selecting file: {str(e)}" def select_all_files(self): """Select all files for deletion""" output_files = self.get_output_gallery() self.selected_for_deletion = set(output_files) selected_names = [os.path.basename(f) for f in self.selected_for_deletion] selection_text = f"Selected ({len(selected_names)}): {', '.join(selected_names) if selected_names else 'None'}" updated_gallery = self.get_output_gallery_with_selection_visual() return updated_gallery, selection_text, f"Selected all {len(selected_names)} files" def clear_file_selection(self): """Clear all file selections""" self.selected_for_deletion = set() updated_gallery = self.get_output_gallery_with_selection_visual() return updated_gallery, "Selected (0): None", "Cleared all selections" def delete_selected_crops(self): """Delete multiple selected crops""" if not self.selected_for_deletion: updated_gallery = self.get_output_gallery_with_selection_visual() return updated_gallery, "No files selected for deletion", "Selected (0): None" try: deleted_count = 0 deleted_names = [] for file_path in list(self.selected_for_deletion): if os.path.exists(file_path): deleted_names.append(os.path.basename(file_path)) os.remove(file_path) deleted_count += 1 # Clear selection after deletion self.selected_for_deletion = set() updated_gallery = self.get_output_gallery_with_selection_visual() return updated_gallery, f"šŸ—‘ļø Deleted {deleted_count} files: {', '.join(deleted_names)}", "Selected (0): None" except Exception as e: updated_gallery = self.get_output_gallery_with_selection_visual() return updated_gallery, f"āŒ Error deleting files: {str(e)}", "Selected (0): None" def download_all_crops(self): """Create a zip file with all cropped images""" if not os.path.exists(self.output_dir) or not os.listdir(self.output_dir): return None, "No crops to download" try: zip_path = os.path.join(self.temp_dir, "cropped_images.zip") with zipfile.ZipFile(zip_path, 'w') as zipf: for filename in os.listdir(self.output_dir): if filename.lower().endswith(('.png', '.jpg', '.jpeg')): file_path = os.path.join(self.output_dir, filename) zipf.write(file_path, filename) return zip_path, f"šŸ“¦ Created zip with {len(os.listdir(self.output_dir))} images" except Exception as e: return None, f"āŒ Error creating zip: {str(e)}" def get_output_gallery_with_selection_visual(self): """Get list of output images with selection indicators using Gradio captions""" if not os.path.exists(self.output_dir): return [] output_files = [] for filename in sorted(os.listdir(self.output_dir)): if filename.lower().endswith(('.png', '.jpg', '.jpeg')): file_path = os.path.join(self.output_dir, filename) # Create clean thumbnail and get dimensions clean_thumb_path, dimensions = self.create_clean_thumbnail(file_path) # Create caption with dimensions and selection status if file_path in self.selected_for_deletion: caption = f"{dimensions} ā€¢ šŸ”“ SELECTED" else: caption = f"{dimensions}" display_item = (clean_thumb_path, caption) output_files.append(display_item) return output_files def refresh_output_gallery(self): """Refresh output gallery with dynamic height""" gallery_data = self.get_output_gallery_with_selection_visual() return gallery_data # Utility Functions def get_image_base_name(self, filename): """Get base name of image file without extension""" return os.path.splitext(filename)[0] def find_caption_file(self, image_path, folder_path): """Find corresponding caption (.txt) file for an image""" base_name = self.get_image_base_name(os.path.basename(image_path)) caption_filename = base_name + ".txt" caption_path = os.path.join(folder_path, caption_filename) # Check if caption file exists if os.path.exists(caption_path): return caption_path # Also check in subdirectories (in case of nested structure) for root, dirs, files in os.walk(folder_path): if caption_filename in files: return os.path.join(root, caption_filename) return None def convert_images_to_rgb(self, folder_path, preserve_captions=True): """Convert grayscale and RGBA images to RGB.""" converted_count = 0 error_count = 0 conversion_log = [] caption_files_preserved = 0 for filename in os.listdir(folder_path): if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff')): file_path = os.path.join(folder_path, filename) try: with Image.open(file_path) as img: original_mode = img.mode # Check if the image needs conversion if img.mode in ['L', 'LA', 'P']: # Grayscale or palette rgb_img = img.convert('RGB') rgb_img.save(file_path, 'JPEG', quality=95) converted_count += 1 conversion_log.append(f"āœ… {filename}: {original_mode} ā†’ RGB") # Check if caption file exists if preserve_captions: caption_path = self.find_caption_file(file_path, folder_path) if caption_path: caption_files_preserved += 1 elif img.mode == 'RGBA': # Create white background for RGBA conversion rgb_img = Image.new('RGB', img.size, (255, 255, 255)) rgb_img.paste(img, mask=img.split()[-1] if len(img.split()) == 4 else None) rgb_img.save(file_path, 'JPEG', quality=95) converted_count += 1 conversion_log.append(f"āœ… {filename}: RGBA ā†’ RGB (white background)") # Check if caption file exists if preserve_captions: caption_path = self.find_caption_file(file_path, folder_path) if caption_path: caption_files_preserved += 1 else: conversion_log.append(f"ā„¹ļø {filename}: Already RGB, skipped") # Still check for caption files if preserve_captions: caption_path = self.find_caption_file(file_path, folder_path) if caption_path: caption_files_preserved += 1 except Exception as e: error_count += 1 conversion_log.append(f"āŒ {filename}: Error - {str(e)}") if preserve_captions and caption_files_preserved > 0: conversion_log.append(f"šŸ“ Caption files preserved: {caption_files_preserved}") return converted_count, error_count, conversion_log def check_and_remove_corrupted_images(self, folder_path, preserve_captions=True): """Check for corrupted/truncated images and move them to corrupted folder.""" corrupted_count = 0 checked_count = 0 corruption_log = [] caption_files_removed = 0 for filename in os.listdir(folder_path): if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff')): file_path = os.path.join(folder_path, filename) checked_count += 1 try: with Image.open(file_path) as img: img.verify() # Verify if the image is corrupted img = Image.open(file_path) # Reopen for further checks ImageOps.exif_transpose(img) # Simple operation to check loadability corruption_log.append(f"āœ… {filename}: OK") except (IOError, UnidentifiedImageError, Exception) as e: corrupted_count += 1 # Find and handle corresponding caption file caption_path = None if preserve_captions: caption_path = self.find_caption_file(file_path, folder_path) # Move corrupted file to corrupted directory corrupted_path = os.path.join(self.corrupted_dir, filename) shutil.move(file_path, corrupted_path) # Move corresponding caption file if it exists if caption_path and os.path.exists(caption_path): caption_filename = os.path.basename(caption_path) corrupted_caption_path = os.path.join(self.corrupted_dir, caption_filename) shutil.move(caption_path, corrupted_caption_path) caption_files_removed += 1 corruption_log.append(f"šŸ—‘ļø {filename}: Corrupted/truncated - moved to quarantine (+ caption file)") else: corruption_log.append(f"šŸ—‘ļø {filename}: Corrupted/truncated - moved to quarantine") if preserve_captions and caption_files_removed > 0: corruption_log.append(f"šŸ“ Caption files also quarantined: {caption_files_removed}") return checked_count, corrupted_count, corruption_log def process_uploaded_dataset(self, zip_file, convert_rgb, check_corruption, preserve_captions): """Process uploaded dataset ZIP file with selected utilities.""" if not zip_file: return None, "No file uploaded", [], "No processing log available" try: # Clear previous utilities processing if os.path.exists(self.processed_dir): shutil.rmtree(self.processed_dir) os.makedirs(self.processed_dir, exist_ok=True) # Extract ZIP file with zipfile.ZipFile(zip_file.name, 'r') as zip_ref: zip_ref.extractall(self.processed_dir) processing_log = [f"šŸ“¦ Extracted ZIP file: {os.path.basename(zip_file.name)}"] # Find image files and caption files in extracted folders image_files = [] caption_files = [] for root, dirs, files in os.walk(self.processed_dir): for file in files: file_path = os.path.join(root, file) if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp')): image_files.append(file_path) elif file.lower().endswith('.txt') and preserve_captions: caption_files.append(file_path) processing_log.append(f"šŸ” Found {len(image_files)} image files") if preserve_captions: processing_log.append(f"šŸ“ Found {len(caption_files)} caption files") if not image_files: return None, "No image files found in ZIP", [], "\n".join(processing_log) # Apply corruption check first (if enabled) if check_corruption: processing_log.append("\nšŸ” CHECKING FOR CORRUPTED IMAGES:") checked, corrupted, corruption_log = self.check_and_remove_corrupted_images(self.processed_dir, preserve_captions) processing_log.extend(corruption_log) processing_log.append(f"šŸ“Š Corruption Check Summary: {checked} checked, {corrupted} corrupted files removed") # Apply RGB conversion (if enabled) if convert_rgb: processing_log.append("\nšŸŽØ CONVERTING TO RGB:") converted, errors, conversion_log = self.convert_images_to_rgb(self.processed_dir, preserve_captions) processing_log.extend(conversion_log) processing_log.append(f"šŸ“Š Conversion Summary: {converted} converted, {errors} errors") # Create new ZIP with processed images and caption files output_zip_path = os.path.join(self.utilities_dir, "processed_dataset.zip") with zipfile.ZipFile(output_zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: # Add all remaining image files for root, dirs, files in os.walk(self.processed_dir): for file in files: file_path = os.path.join(root, file) # Get relative path for ZIP arcname = os.path.relpath(file_path, self.processed_dir) # Add image files if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp')): zipf.write(file_path, arcname) # Add caption files if preservation is enabled elif file.lower().endswith('.txt') and preserve_captions: zipf.write(file_path, arcname) # Get final file counts final_image_count = 0 final_caption_count = 0 for root, dirs, files in os.walk(self.processed_dir): for file in files: if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp')): final_image_count += 1 elif file.lower().endswith('.txt') and preserve_captions: final_caption_count += 1 if preserve_captions and final_caption_count > 0: processing_log.append(f"\nāœ… Processing complete! Final dataset contains {final_image_count} images and {final_caption_count} caption files") else: processing_log.append(f"\nāœ… Processing complete! Final dataset contains {final_image_count} images") # Create gallery of processed images (first 20 for preview) preview_images = [] count = 0 for root, dirs, files in os.walk(self.processed_dir): for file in sorted(files): if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff')) and count < 20: file_path = os.path.join(root, file) preview_images.append(file_path) count += 1 if count >= 20: break if count >= 20: break status_message = f"āœ… Processing complete! {final_image_count} images" if preserve_captions and final_caption_count > 0: status_message += f" and {final_caption_count} caption files" status_message += " ready for download" return output_zip_path, status_message, preview_images, "\n".join(processing_log) except Exception as e: return None, f"āŒ Error processing dataset: {str(e)}", [], f"Error: {str(e)}" def download_all_crops_with_utilities(self, convert_rgb, check_corruption): """Create a zip file with all cropped images, optionally processed through utilities""" if not os.path.exists(self.output_dir) or not os.listdir(self.output_dir): return None, "No crops to download" try: processing_log = [] # If utilities are requested, process the crops first if convert_rgb or check_corruption: # Copy crops to utilities processing folder temp_process_dir = os.path.join(self.utilities_dir, "temp_crops") if os.path.exists(temp_process_dir): shutil.rmtree(temp_process_dir) os.makedirs(temp_process_dir, exist_ok=True) # Copy all crops to temp processing directory for filename in os.listdir(self.output_dir): if filename.lower().endswith(('.png', '.jpg', '.jpeg')): src_path = os.path.join(self.output_dir, filename) dst_path = os.path.join(temp_process_dir, filename) shutil.copy2(src_path, dst_path) processing_log.append(f"šŸ“¦ Copied {len(os.listdir(temp_process_dir))} crops for processing") # Apply corruption check first (if enabled) - no caption preservation for crops if check_corruption: processing_log.append("šŸ” Checking for corrupted crops...") checked, corrupted, corruption_log = self.check_and_remove_corrupted_images(temp_process_dir, False) processing_log.append(f"Corruption check: {checked} checked, {corrupted} corrupted") # Apply RGB conversion (if enabled) - no caption preservation for crops if convert_rgb: processing_log.append("šŸŽØ Converting crops to RGB...") converted, errors, conversion_log = self.convert_images_to_rgb(temp_process_dir, False) processing_log.append(f"RGB conversion: {converted} converted, {errors} errors") # Create ZIP from processed crops zip_path = os.path.join(self.temp_dir, "processed_cropped_images.zip") source_dir = temp_process_dir processing_status = " (Processed)" else: # Create ZIP from original crops zip_path = os.path.join(self.temp_dir, "cropped_images.zip") source_dir = self.output_dir processing_status = "" with zipfile.ZipFile(zip_path, 'w') as zipf: for filename in os.listdir(source_dir): if filename.lower().endswith(('.png', '.jpg', '.jpeg')): file_path = os.path.join(source_dir, filename) zipf.write(file_path, filename) final_count = len([f for f in os.listdir(source_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]) status_message = f"šŸ“¦ Created zip with {final_count} images{processing_status}" if processing_log: status_message += f"\n\nProcessing Log:\n" + "\n".join(processing_log) return zip_path, status_message except Exception as e: return None, f"āŒ Error creating zip: {str(e)}" # Initialize the app app = ImagePrepApp() # Create the Gradio interface with tabs with gr.Blocks(title="Image Prep Tool", theme=gr.themes.Soft(), css=""" #main_crop_image { max-height: none !important; height: auto !important; } #main_crop_image img { max-height: none !important; height: auto !important; max-width: 100% !important; } /* Remove Gradio's default blue selection border */ #output_gallery .selected, #output_gallery .thumbnail.selected { border: 1px solid #374151 !important; /* Use default border instead of blue */ box-shadow: none !important; transform: none !important; } /* Hover effect for better UX */ #output_gallery .thumbnail:hover { border: 2px solid #4a9eff !important; box-shadow: 0 0 8px rgba(74, 158, 255, 0.5) !important; transform: scale(1.02) !important; transition: all 0.2s ease !important; } .about-section { background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; padding: 20px; border-radius: 10px; margin-bottom: 20px; } .step-card { background: #7382bb; border-left: 4px solid #007bff; padding: 15px; margin: 10px 0; border-radius: 0 8px 8px 0; color: #212529; box-shadow: 0 2px 4px rgba(0,0,0,0.1); } /* Remove drawer styling - keep it simple */ .gallery-drawer { border: 2px solid #e0e0e0; border-radius: 10px; padding: 15px; margin-bottom: 20px; background: #fafafa; transition: all 0.3s ease; } /* Drawer toggle button styling */ .drawer-toggle { background: linear-gradient(135deg, #4CAF50 0%, #45a049 100%) !important; color: white !important; font-weight: bold !important; border: none !important; border-radius: 8px !important; padding: 10px 15px !important; margin-bottom: 10px !important; transition: all 0.3s ease !important; box-shadow: 0 2px 4px rgba(0,0,0,0.2) !important; } .drawer-toggle:hover { background: linear-gradient(135deg, #45a049 0%, #3d8b40 100%) !important; transform: translateY(-2px) !important; box-shadow: 0 4px 8px rgba(0,0,0,0.3) !important; } """) as demo: gr.Markdown("# āœ‚ļø PixelPruner") gr.Markdown("Quick and Easy Dataset Prep!") with gr.Tabs() as tabs: # TAB 0: ABOUT & USAGE (NEW - FRONT TAB) with gr.Tab("šŸ“– About & Usage", id="about_tab"): with gr.Column(): # Hero Section gr.HTML("""

šŸŽÆ Welcome to PixelPruner v1.2

Streamline your workflow and achieve perfect image crops every time with PixelPruner!

""") # How to Use gr.Markdown("## šŸš€ How to Use PixelPruner") gr.HTML("""

Step 1: Load Your Images

Go to the šŸ“‚ Load Images tab and upload your image files. PixelPruner supports PNG, JPG, JPEG, GIF, BMP, and TIFF formats.

""") gr.HTML("""

Step 2: Select and Configure

Move to the āœ‚ļø Crop Images tab. Click on thumbnails to select images, then choose your desired crop dimensions.

""") gr.HTML("""

Step 3: Crop with Precision

Click anywhere on the main image to crop at that location. Use the zoom slider to control crop area size.

""") gr.HTML("""

Step 4: Save and Manage

Save crops you like, then use the šŸ“¦ Preview & Download Crops tab to review, delete unwanted crops, and download your final dataset as a ZIP file.

""") # Technical Details gr.Markdown("## āš™ļø Technical Details") with gr.Row(): with gr.Column(): gr.Markdown(""" **šŸ–¼ļø Image Processing:** - High-quality LANCZOS resampling for all operations - Crops are extracted from original full-resolution images - Display images are optimized for performance without quality loss - Supports images of any size and aspect ratio """) with gr.Column(): gr.Markdown(""" **šŸ“ File Management:** - Temporary files are automatically cleaned up - Crops saved as PNG for maximum quality - Batch operations for efficient workflow - ZIP export maintains original quality """) gr.Markdown("---") gr.Markdown("*Ready to start cropping? Head over to the **šŸ“‚ Load Images** tab to begin!*") # TAB 1: IMAGE LOADING with gr.Tab("šŸ“‚ Load Images", id="load_tab"): gr.Markdown("## Upload Your Images") gr.Markdown("Select multiple image files to begin processing") with gr.Row(): with gr.Column(): files_input = gr.File( label="Choose Image Files", file_count="multiple", file_types=["image"], height=200 ) load_status = gr.Textbox( label="Load Status", interactive=False, placeholder="No images loaded yet..." ) gr.Markdown("### šŸ“‹ Supported Formats") gr.Markdown("PNG, JPG, JPEG, GIF, BMP, TIFF") # TAB 2: CROPPING with gr.Tab("āœ‚ļø Crop Images", id="crop_tab", visible=False) as crop_tab: gr.Markdown("## Image Cropping") with gr.Row(): # Left column - Image selection and cropping with gr.Column(scale=2): # Collapsible Gallery Section gallery_visible = gr.State(True) # Track drawer visibility state with gr.Row(): gr.Markdown("### šŸ–¼ļø Select Image to Crop") toggle_btn = gr.Button( "šŸ”½ Hide Gallery", variant="secondary", size="sm", scale=1 ) with gr.Column(visible=True) as gallery_section: image_gallery = gr.Gallery( label="Click thumbnail to select image", show_label=True, columns=8, object_fit="cover", allow_preview=False, container=True, interactive=True ) with gr.Row(): image_info = gr.Textbox( label="Current Image Info", interactive=False, placeholder="Select an image from the gallery above", scale=2 ) display_size_dropdown = gr.Dropdown( choices=list(app.display_sizes.keys()), value="Medium (800x600)", label="Display Size", scale=1 ) gr.Markdown("### šŸŽÆ Click on image to crop at that location") current_image = gr.Image( label="Click to Crop", type="pil", interactive=True, height=600, container=True, show_label=True, elem_id="main_crop_image", sources=[] ) # Navigation buttons with gr.Row(): prev_btn = gr.Button("ā¬…ļø Previous Image", variant="secondary", scale=1) next_btn = gr.Button("Next Image āž”ļø", variant="secondary", scale=1) # Right column - Crop settings and preview with gr.Column(scale=1): gr.Markdown("### šŸ”² Crop Settings") crop_preset = gr.Dropdown( choices=list(app.crop_presets.keys()), value="512x512", label="Crop Dimensions" ) with gr.Group(visible=False) as custom_group: custom_width = gr.Number(label="Custom Width", value=512, minimum=1) custom_height = gr.Number(label="Custom Height", value=512, minimum=1) # Zoom slider gr.Markdown("### šŸ” Zoom Control") gr.Markdown("*Higher zoom = closer view (smaller crop area)*") zoom_slider = gr.Slider( minimum=0.1, maximum=3.0, step=0.1, value=1.0, label="Zoom Level", interactive=True ) zoom_info = gr.Textbox( value="Zoom: 1.0x", label="Zoom Info", interactive=False ) # Hidden fields to store current crop dimensions current_crop_width = gr.Number(value=512, visible=False) current_crop_height = gr.Number(value=512, visible=False) gr.Markdown("### šŸ–¼ļø Crop Preview") cropped_image = gr.Image( label="Cropped Result", type="pil", height=300, sources=[] ) crop_status = gr.Textbox( label="Crop Info", interactive=False, placeholder="Click on image to create crop" ) save_btn = gr.Button("šŸ’¾ Save Crop", variant="primary", size="lg") save_status = gr.Textbox(label="Save Status", interactive=False) # TAB 3: DOWNLOAD MANAGEMENT with gr.Tab("šŸ“¦ Preview & Download Crops", id="download_tab", visible=False) as download_tab: gr.Markdown("## Manage Your Cropped Images") with gr.Row(): refresh_btn = gr.Button("šŸ”„ Refresh Gallery", variant="secondary") download_all_btn = gr.Button("šŸ“¦ Prepare ZIP for Download", variant="primary") download_status = gr.Textbox(label="Download Status", interactive=False) download_file = gr.File(label="Download ZIP", visible=False) # Processing options for download (initially hidden) with gr.Group(visible=False) as processing_options: gr.Markdown("### šŸ”§ Apply Utilities Before Download?") with gr.Row(): download_convert_rgb = gr.Checkbox(label="šŸŽØ Convert to RGB", value=False) download_corruption_check = gr.Checkbox(label="šŸ” Check for Corruption", value=False) with gr.Row(): download_confirm_btn = gr.Button("šŸ“¦ Create ZIP", variant="primary") download_skip_btn = gr.Button("ā­ļø Skip Utilities", variant="secondary") gr.Markdown("### šŸ–¼ļø Saved Crops") gr.Markdown("**Left-click on images to select/deselect for deletion**") output_gallery = gr.Gallery( label="Cropped Images", show_label=True, columns=6, # Reduced columns to give more space object_fit="contain", # Changed from "cover" to "contain" to preserve aspect ratio allow_preview=False, show_download_button=False, selected_index=None, # Prevent persistent selection elem_id="output_gallery" ) # Multi-select for deletion gr.Markdown("### šŸ—‘ļø Batch Delete") gr.Markdown("*Click images above to select them, then use buttons below*") with gr.Row(): select_all_btn = gr.Button("ā˜‘ļø Select All", variant="secondary") clear_selection_btn = gr.Button("šŸ”„ Clear Selection", variant="secondary") delete_selected_btn = gr.Button("šŸ—‘ļø Delete Selected", variant="stop") selected_files_display = gr.Textbox( label="Selected Files (click images to select)", interactive=False, placeholder="No files selected - click on images above to select them" ) delete_status = gr.Textbox(label="Delete Status", interactive=False) # TAB 4: UTILITIES with gr.Tab("šŸ”§ Utilities", id="utilities_tab"): gr.Markdown("## Dataset Processing Utilities") gr.Markdown("Upload a ZIP file containing images to process with various utilities") with gr.Row(): with gr.Column(scale=1): gr.Markdown("### šŸ“¤ Upload Dataset") dataset_zip = gr.File( label="Upload ZIP Dataset", file_types=[".zip"], height=150 ) gr.Markdown("### āš™ļø Processing Options") with gr.Group(): convert_rgb_check = gr.Checkbox( label="šŸŽØ Convert to RGB", value=False, info="Convert grayscale, RGBA, and palette images to RGB format" ) corruption_check = gr.Checkbox( label="šŸ” Remove Corrupted Images", value=True, info="Scan for and remove corrupted/truncated image files" ) preserve_captions_check = gr.Checkbox( label="šŸ“ Preserve Caption Files", value=True, info="Keep .txt caption files paired with images (for LoRA training)" ) process_btn = gr.Button("šŸš€ Process Dataset", variant="primary", size="lg") with gr.Row(): utilities_status = gr.Textbox( label="Processing Status", interactive=False, placeholder="Upload a ZIP file and click Process Dataset..." ) utilities_download = gr.File( label="Download Processed Dataset", visible=False ) with gr.Column(scale=1): gr.Markdown("### šŸ“‹ Processing Log") processing_log = gr.Textbox( label="Detailed Log", lines=15, max_lines=20, interactive=False, placeholder="Processing log will appear here..." ) gr.Markdown("### šŸ–¼ļø Preview of Processed Images") utilities_gallery = gr.Gallery( label="Processed Images Preview (First 20)", show_label=True, columns=6, object_fit="contain", allow_preview=True, height=300 ) gr.Markdown("---") gr.Markdown("### ā„¹ļø Utility Information - Dataset processing scripts by Jomcey!") with gr.Row(): with gr.Column(): gr.Markdown(""" **šŸŽØ Convert to RGB:** - Converts grayscale images to RGB format - Converts RGBA images to RGB with white background - Converts palette images to RGB - Saves as JPEG with 95% quality """) with gr.Column(): gr.Markdown(""" **šŸ” Remove Corrupted Images:** - Scans for truncated or corrupted image files - Removes files that can't be properly loaded - Quarantines problematic files separately - Ensures dataset integrity """) with gr.Column(): gr.Markdown(""" **šŸ“ Preserve Caption Files:** - Keeps .txt files paired with images - Essential for LoRA training datasets - Removes captions only when images are corrupted - Maintains proper image-caption relationships """) # Event handlers # Toggle gallery section visibility def update_toggle_button_and_gallery(current_visibility): new_visibility = not current_visibility if new_visibility: button_text = "šŸ”½ Hide Gallery" else: button_text = "šŸ”¼ Show Gallery" return new_visibility, gr.update(value=button_text), gr.update(visible=new_visibility) toggle_btn.click( update_toggle_button_and_gallery, inputs=[gallery_visible], outputs=[gallery_visible, toggle_btn, gallery_section] ) # Load images files_input.upload( app.load_images_from_folder, inputs=[files_input], outputs=[load_status, image_gallery, crop_tab] ) # Gallery selection in crop tab image_gallery.select( app.select_from_gallery, inputs=[display_size_dropdown], outputs=[current_image, image_info, zoom_slider] ) # Display size change display_size_dropdown.change( app.update_display_size, inputs=[display_size_dropdown], outputs=[current_image, image_info] ) # Navigation buttons prev_btn.click( lambda display_size: app.navigate_image("prev", display_size), inputs=[display_size_dropdown], outputs=[current_image, image_info, zoom_slider] ) next_btn.click( lambda display_size: app.navigate_image("next", display_size), inputs=[display_size_dropdown], outputs=[current_image, image_info, zoom_slider] ) # Crop preset change handler crop_preset.change( app.update_crop_dimensions, inputs=[crop_preset, current_crop_width, current_crop_height], outputs=[custom_group, current_crop_width, current_crop_height] ) # Custom dimensions update custom_width.change( lambda w: w, inputs=[custom_width], outputs=[current_crop_width] ) custom_height.change( lambda h: h, inputs=[custom_height], outputs=[current_crop_height] ) # Zoom slider update zoom_slider.change( app.update_zoom, inputs=[zoom_slider], outputs=[zoom_info] ) # Image click for cropping current_image.select( app.process_crop_click, inputs=[current_image, crop_preset, custom_width, custom_height, zoom_slider, display_size_dropdown], outputs=[cropped_image, crop_status] ) # Save crop save_btn.click( app.save_crop, inputs=[cropped_image], outputs=[save_status, download_tab] ) # Download tab functions refresh_btn.click( app.refresh_output_gallery, outputs=[output_gallery] ) # Auto-refresh gallery when download tab is selected download_tab.select( app.refresh_output_gallery, outputs=[output_gallery] ) # Gallery selection for deletion output_gallery.select( app.toggle_file_selection, outputs=[output_gallery, selected_files_display, delete_status] ) # Batch selection controls select_all_btn.click( app.select_all_files, outputs=[output_gallery, selected_files_display, delete_status] ) clear_selection_btn.click( app.clear_file_selection, outputs=[output_gallery, selected_files_display, delete_status] ) delete_selected_btn.click( app.delete_selected_crops, outputs=[output_gallery, delete_status, selected_files_display] ) # Download all crops - show processing options download_all_btn.click( lambda: gr.update(visible=True), outputs=[processing_options] ) # Confirm download with utilities download_confirm_btn.click( lambda rgb, corruption: app.download_all_crops_with_utilities(rgb, corruption), inputs=[download_convert_rgb, download_corruption_check], outputs=[download_file, download_status] ).then( lambda: (gr.update(visible=True), gr.update(visible=False)), outputs=[download_file, processing_options] ) # Skip utilities and download directly download_skip_btn.click( lambda: app.download_all_crops_with_utilities(False, False), outputs=[download_file, download_status] ).then( lambda: (gr.update(visible=True), gr.update(visible=False)), outputs=[download_file, processing_options] ) # Utilities tab functions process_btn.click( app.process_uploaded_dataset, inputs=[dataset_zip, convert_rgb_check, corruption_check, preserve_captions_check], outputs=[utilities_download, utilities_status, utilities_gallery, processing_log] ).then( lambda x: gr.update(visible=True) if x else gr.update(visible=False), inputs=[utilities_download], outputs=[utilities_download] ) if __name__ == "__main__": demo.launch(share=True)