import os import zipfile import shutil import time from PIL import Image, ImageDraw import io from rembg import remove import gradio as gr from concurrent.futures import ThreadPoolExecutor from diffusers import StableDiffusionPipeline from transformers import pipeline import numpy as np import json import torch import logging # Load Stable Diffusion Model def load_stable_diffusion_model(): device = "cuda" if torch.cuda.is_available() else "cpu" pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float32).to(device) return pipe # Initialize the model globally sd_model = load_stable_diffusion_model() def remove_background_rembg(input_path): print(f"Removing background using rembg for image: {input_path}") with open(input_path, 'rb') as i: input_image = i.read() output_image = remove(input_image) img = Image.open(io.BytesIO(output_image)).convert("RGBA") return img def remove_background_bria(input_path): """Remove background using the Bria model.""" print("Removing background using Bria for the image.") # Create a segmentation pipeline pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True) # Load the image input_image = Image.open(input_path).convert("RGBA") # Get the segmentation output pillow_mask = pipe(input_image, return_mask=True) # Outputs a pillow mask print("Mask obtained:", pillow_mask) # Debugging output # Create an output image based on the mask output_image = Image.new("RGBA", input_image.size) # Use the mask to create the output image for x in range(input_image.width): for y in range(input_image.height): # Assuming mask is in a binary format where foreground is True if pillow_mask.getpixel((x, y)) > 0: # Adjust based on actual mask values output_image.putpixel((x, y), input_image.getpixel((x, y))) else: output_image.putpixel((x, y), (0, 0, 0, 0)) # Set to transparent return output_image # Function to process images using prompts def text_to_image(prompt): os.makedirs("generated_images", exist_ok=True) # Ensure the directory exists image = sd_model(prompt).images[0] # Generate image using the model # Create a sanitized filename by replacing spaces with underscores image_path = f"generated_images/{prompt.replace(' ', '_')}.png" image.save(image_path) # Save the generated image return image, image_path # Return the image and its path # Function to modify an image based on a text prompt def text_image_to_image(input_image, prompt): os.makedirs("generated_images", exist_ok=True) # Ensure the directory exists # Convert input image to PIL Image if necessary if not isinstance(input_image, Image.Image): input_image = Image.open(input_image) # Load image from path if given as string # Generate modified image using the model with the input image and prompt modified_image = sd_model(prompt, init_image=input_image, strength=0.75).images[0] # Create a sanitized filename for the modified image image_path = f"generated_images/{prompt.replace(' ', '_')}_modified.png" modified_image.save(image_path) # Save the modified image return modified_image, image_path # Return the modified image and its path def get_bounding_box_with_threshold(image, threshold): # Convert image to numpy array img_array = np.array(image) # Get alpha channel alpha = img_array[:, :, 3] # Find rows and columns where alpha > threshold rows = np.any(alpha > threshold, axis=1) cols = np.any(alpha > threshold, axis=0) # Find the bounding box if np.any(rows) and np.any(cols): top, bottom = np.where(rows)[0][[0, -1]] left, right = np.where(cols)[0][[0, -1]] return (left, top, right, bottom) else: return None def position_logic(image_path, canvas_size, padding_top, padding_right, padding_bottom, padding_left, use_threshold=True): """Position and resize an image based on cropping and padding requirements.""" image = Image.open(image_path).convert("RGBA") # Get the bounding box of the non-blank area with threshold bbox = get_bounding_box_with_threshold(image, threshold=10) if use_threshold else image.getbbox() log = [] if bbox: width, height = image.size cropped_sides = [] tolerance = 30 # Define a constant for transparency tolerance # Check each edge for non-transparent pixels for edge in ['top', 'bottom', 'left', 'right']: if edge == 'top' and any(image.getpixel((x, 0))[3] > tolerance for x in range(width)): cropped_sides.append(edge) elif edge == 'bottom' and any(image.getpixel((x, height - 1))[3] > tolerance for x in range(width)): cropped_sides.append(edge) elif edge == 'left' and any(image.getpixel((0, y))[3] > tolerance for y in range(height)): cropped_sides.append(edge) elif edge == 'right' and any(image.getpixel((width - 1, y))[3] > tolerance for y in range(height)): cropped_sides.append(edge) # Log cropping information info_message = f"Info for {os.path.basename(image_path)}: The following sides of the image may contain cropped objects: {', '.join(cropped_sides) if cropped_sides else 'none'}" print(info_message) log.append({"info": info_message}) # Crop the image to the bounding box image = image.crop(bbox) log.append({"action": "crop", "bbox": [str(bbox[0]), str(bbox[1]), str(bbox[2]), str(bbox[3])]}) # Calculate new size to expand the image target_width, target_height = canvas_size aspect_ratio = image.width / image.height # Handling image positioning and resizing based on cropped sides if len(cropped_sides) == 4: # Center crop if cropped on all sides if aspect_ratio > 1: # Landscape new_height = target_height new_width = int(new_height * aspect_ratio) left = (new_width - target_width) // 2 image = image.resize((new_width, new_height), Image.LANCZOS).crop((left, 0, left + target_width, target_height)) else: # Portrait or square new_width = target_width new_height = int(new_width / aspect_ratio) top = (new_height - target_height) // 2 image = image.resize((new_width, new_height), Image.LANCZOS).crop((0, top, target_width, top + target_height)) log.append({"action": "center_crop_resize", "new_size": f"{target_width}x{target_height}"}) x, y = 0, 0 elif not cropped_sides: # Expand from center new_height = target_height - padding_top - padding_bottom new_width = int(new_height * aspect_ratio) if new_width > target_width - padding_left - padding_right: new_width = target_width - padding_left - padding_right new_height = int(new_width / aspect_ratio) image = image.resize((new_width, new_height), Image.LANCZOS) log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)}) x = (target_width - new_width) // 2 y = target_height - new_height - padding_bottom else: # Handle specific cropped side cases new_width, new_height = target_width, target_height # Default values for resizing for side in cropped_sides: if side in ["top", "bottom"]: new_height = target_height - (padding_top + padding_bottom) if side in ["left", "right"]: new_width = target_width - (padding_left + padding_right) image = image.resize((new_width, new_height), Image.LANCZOS) log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)}) # Determine position based on cropped sides x = (target_width - new_width) // 2 if "left" not in cropped_sides and "right" not in cropped_sides else (0 if "left" in cropped_sides else target_width - new_width) y = (0 if "top" in cropped_sides else target_height - new_height) log.append({"action": "position", "x": str(x), "y": str(y)}) return log, image, x, y # Constants for canvas sizes and paddings CANVAS_SIZES = { 'Rox': ((1080, 1080), (112, 125, 116, 125)), 'Columbia': ((730, 610), (30, 105, 35, 105)), 'Zalora': ((763, 1100), (50, 50, 200, 50)) } def process_single_image(image_path, output_folder, bg_method, canvas_size_name, output_format, bg_choice, custom_color, watermark_path=None): """Processes a single image by removing its background and applying various transformations. Args: image_path (str): Path to the input image. output_folder (str): Path to the output folder. bg_method (str): Background removal method ('rembg' or 'bria'). canvas_size_name (str): Name of the canvas size. output_format (str): Desired output format ('JPG' or 'PNG'). bg_choice (str): Background choice ('white', 'custom', or 'transparent'). custom_color (tuple): Custom background color as an RGBA tuple. watermark_path (str, optional): Path to a watermark image. Returns: tuple: A tuple containing the output path and log of actions performed. """ try: canvas_size, (padding_top, padding_right, padding_bottom, padding_left) = CANVAS_SIZES[canvas_size_name] filename = os.path.basename(image_path) logging.info(f"Processing image: {filename}") # Remove background if bg_method == 'rembg': image_with_no_bg = remove_background_rembg(image_path) elif bg_method == 'bria': image_with_no_bg = remove_background_bria(image_path) else: raise ValueError("Invalid background method specified.") # Temporary file for processed image temp_image_path = os.path.join(output_folder, f"temp_{filename}") image_with_no_bg.save(temp_image_path, format='PNG') log, new_image, x, y = position_logic(temp_image_path, canvas_size, padding_top, padding_right, padding_bottom, padding_left) # Create canvas based on background choice if bg_choice == 'white': canvas = Image.new("RGBA", canvas_size, "WHITE") elif bg_choice == 'custom': canvas = Image.new("RGBA", canvas_size, custom_color) else: # transparent canvas = Image.new("RGBA", canvas_size, (0, 0, 0, 0)) # Paste the resized image onto the canvas canvas.paste(new_image, (x, y), new_image) log.append({"action": "paste", "position": [str(x), str(y)]}) # Add watermark if applicable if watermark_path: watermark = Image.open(watermark_path).convert("RGBA") # Hitung posisi untuk menempatkan watermark di tengah watermark_width, watermark_height = watermark.size canvas_width, canvas_height = canvas.size # Hitung posisi (x, y) untuk watermark x = (canvas_width - watermark_width) // 2 y = (canvas_height - watermark_height) // 2 # Tempelkan watermark canvas.paste(watermark, (x, y), watermark) log.append({"action": "add_watermark"}) # Determine output format output_ext = 'jpg' if output_format == 'JPG' else 'png' output_filename = f"{os.path.splitext(filename)[0]}.{output_ext}" output_path = os.path.join(output_folder, output_filename) # Save the canvas in the desired format if output_format == 'JPG': canvas.convert('RGB').save(output_path, format='JPEG') else: canvas.save(output_path, format='PNG') # Clean up temporary file os.remove(temp_image_path) logging.info(f"Processed image path: {output_path}") return [(output_path, image_path)], log except Exception as e: logging.error(f"Error processing {filename}: {e}") return None, None # Set up logging logging.basicConfig(level=logging.INFO) def process_images(input_files, bg_method='rembg', watermark_path=None, canvas_size='Rox', output_format='PNG', bg_choice='transparent', custom_color="#ffffff", num_workers=4, progress=gr.Progress()): """Processes images by removing backgrounds and applying various transformations. Args: input_files (str or list): Path to a ZIP file or a list of image paths. bg_method (str): Background removal method ('rembg' or 'bria'). watermark_path (str, optional): Path to a watermark image. canvas_size (str): Name of the canvas size. output_format (str): Desired output format ('PNG' or 'JPG'). bg_choice (str): Background choice ('transparent', 'white', or 'custom'). custom_color (str): Custom background color in hex format. num_workers (int): Number of parallel workers for processing. progress (gr.Progress): Progress tracking interface. Returns: tuple: A tuple containing original images, processed images, output zip path, and total processing time. """ start_time = time.time() output_folder = "processed_images" if os.path.exists(output_folder): shutil.rmtree(output_folder) os.makedirs(output_folder) processed_images = [] original_images = [] all_logs = [] image_files = [] if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')): # Handle zip file input_folder = "temp_input" if os.path.exists(input_folder): shutil.rmtree(input_folder) os.makedirs(input_folder) try: with zipfile.ZipFile(input_files, 'r') as zip_ref: zip_ref.extractall(input_folder) image_files = [os.path.join(input_folder, f) for f in os.listdir(input_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp'))] except zipfile.BadZipFile as e: logging.error(f"Error extracting zip file: {e}") return [], None, 0 elif isinstance(input_files, list): # Handle multiple files image_files = input_files else: # Handle single file image_files = [input_files] total_images = len(image_files) logging.info(f"Total images to process: {total_images}") avg_processing_time = 0 with ThreadPoolExecutor(max_workers=num_workers) as executor: future_to_image = {executor.submit(process_single_image, image_path, output_folder, bg_method, canvas_size, output_format, bg_choice, custom_color, watermark_path): image_path for image_path in image_files} for idx, future in enumerate(future_to_image): try: start_time_image = time.time() result, log = future.result() end_time_image = time.time() image_processing_time = end_time_image - start_time_image # Update average processing time avg_processing_time = (avg_processing_time * idx + image_processing_time) / (idx + 1) if result: processed_images.extend(result) original_images.append(future_to_image[future]) all_logs.append({os.path.basename(future_to_image[future]): log}) # Estimate remaining time remaining_images = total_images - (idx + 1) estimated_remaining_time = remaining_images * avg_processing_time progress((idx + 1) / total_images, f"{idx + 1}/{total_images} images processed. Estimated time remaining: {estimated_remaining_time:.2f} seconds") except Exception as e: logging.error(f"Error processing image {future_to_image[future]}: {e}") output_zip_path = "processed_images.zip" with zipfile.ZipFile(output_zip_path, 'w') as zipf: for file, _ in processed_images: zipf.write(file, os.path.basename(file)) # Write the comprehensive log for all images with open(os.path.join(output_folder, 'process_log.json'), 'w') as log_file: json.dump(all_logs, log_file, indent=4) logging.info("Comprehensive log saved to %s", os.path.join(output_folder, 'process_log.json')) end_time = time.time() processing_time = end_time - start_time logging.info(f"Processing time: {processing_time} seconds") return original_images, processed_images, output_zip_path, processing_time # Set up logging logging.basicConfig(level=logging.INFO) def gradio_interface(input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers): """Handles input files and processes them accordingly.""" progress = gr.Progress() watermark_path = watermark.name if watermark else None # Check the type of input_files and call the process_images function if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')): return process_images(input_files, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress) elif isinstance(input_files, list): return process_images(input_files, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress) else: return process_images(input_files.name, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress) def show_color_picker(bg_choice): """Shows the color picker if 'custom' background is selected.""" return gr.update(visible=(bg_choice == 'custom')) def update_compare(evt: gr.SelectData): """Updates the displayed images and their ratios when a processed image is selected.""" if isinstance(evt.value, dict) and 'caption' in evt.value: input_path = evt.value['caption'].split("Input: ")[-1] output_path = evt.value['image']['path'] # Open the original and processed images original_img = Image.open(input_path) processed_img = Image.open(output_path) # Calculate the aspect ratios original_ratio = f"{original_img.width}x{original_img.height}" processed_ratio = f"{processed_img.width}x{processed_img.height}" return (gr.update(value=input_path), gr.update(value=output_path), gr.update(value=original_ratio), gr.update(value=processed_ratio) ) else: logging.warning("No caption found in selection") return (gr.update(value=None),) * 4 def process(input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers): """Processes the images and returns the results.""" try: _, processed_images, zip_path, time_taken = gradio_interface(input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers) processed_images_with_captions = [(img, f"Input: {caption}") for img, caption in processed_images] return processed_images_with_captions, zip_path, f"{time_taken:.2f} seconds" except Exception as e: logging.error(f"Error in processing images: {e}") return [], None, "Error during processing" with gr.Blocks(theme="NoCrypt/miku@1.2.2") as iface: gr.Markdown("# 🎨 Creative Image Suite: Generate, Modify, and Enhance Your Visuals") gr.Markdown(""" **Unlock your creativity with our comprehensive image processing tool! This suite offers three powerful features:** 1. **✏️ Text to Image**: Transform your ideas into stunning visuals by simply entering a descriptive text prompt. Watch your imagination come to life! 2. **🖼️ Image to Image**: Enhance existing images by providing a text description of the modifications you want. Upload any image and specify the changes as you wish to create a unique masterpiece. 3. **🖌️ Image Background Removal and Resizing**: Effortlessly remove backgrounds from images, resize them, and even add watermarks (opitonal). Upload single images or zip files, choose your desired settings, and let our tool process everything seamlessly. """) # Fitur Text to Image gr.Markdown("## Text to Image Feature") gr.Markdown(""" **Example Prompts:** - *A serene mountain landscape at sunset.* - *A futuristic city skyline with flying cars.* - *A whimsical forest filled with colorful mushrooms and fairies.* - *A close-up of a vibrant butterfly resting on a flower.* This feature allows you to create a new image based on a text description. Simply enter your idea in a sentence, and the system will generate an image that matches it. """) gr.Markdown("### ⚠️ Note:") gr.Markdown("Processing may take a while due to the free CPU resources on Hugging Face Spaces. Please be patient!") with gr.Row(): prompt_input = gr.Textbox(label="Enter your prompt for image generation:") generate_button = gr.Button("Generate Image") output_image = gr.Image(label="Generated Image") download_button = gr.File(label="Download Generated Image", type="filepath") generate_button.click(text_to_image, inputs=prompt_input, outputs=[output_image, download_button]) # Fitur Text Image to Image gr.Markdown("## Image to Image Feature") gr.Markdown(""" **Example Prompts:** - *Change the sky to a starry night with a full moon.* - *Add a rainbow across the horizon in this beach scene.* - *Make the flowers in the garden bloom in shades of blue.* - *Transform the cat's fur to a bright orange color.* This feature lets you modify an existing image by adding a text description. Upload an image, specify what you want to change, and the system will alter the image accordingly. """) gr.Markdown("### ⚠️ Note:") gr.Markdown("Processing may take a while due to the free CPU resources on Hugging Face Spaces. Please be patient!") with gr.Row(): input_image = gr.Image(label="Upload Image for Modification", type="pil") prompt_modification = gr.Textbox(label="Enter your prompt for modification:") modify_button = gr.Button("Modify Image") modified_output_image = gr.Image(label="Modified Image") download_modified_button = gr.File(label="Download Modified Image", type="filepath") modify_button.click(text_image_to_image, inputs=[input_image, prompt_modification], outputs=[modified_output_image, download_modified_button]) gr.Markdown("## Image Background Removal and Resizing with Optional Watermark") gr.Markdown("Choose to upload multiple images or a ZIP/RAR file, select the crop mode, optionally upload a watermark image, and choose the output format.") with gr.Row(): input_files = gr.File(label="Upload Image or ZIP/RAR file", file_types=[".zip", ".rar", "image"], interactive=True) watermark = gr.File(label="Upload Watermark Image (Optional)", file_types=[".png"]) with gr.Row(): canvas_size = gr.Radio(choices=["Rox", "Columbia", "Zalora"], label="Canvas Size", value="Rox") output_format = gr.Radio(choices=["PNG", "JPG"], label="Output Format", value="JPG") num_workers = gr.Slider(minimum=1, maximum=16, step=1, label="Number of Workers", value=5) with gr.Row(): bg_method = gr.Radio(choices=["bria", "rembg"], label="Background Removal Method", value="bria") bg_choice = gr.Radio(choices=["transparent", "white", "custom"], label="Background Choice", value="white") custom_color = gr.ColorPicker(label="Custom Background Color", value="#ffffff", visible=False) process_button = gr.Button("Process Images") with gr.Row(): gallery_processed = gr.Gallery(label="Processed Images") with gr.Row(): image_original = gr.Image(label="Original Images", interactive=False) image_processed = gr.Image(label="Processed Images", interactive=False) with gr.Row(): original_ratio = gr.Textbox(label="Original Ratio") processed_ratio = gr.Textbox(label="Processed Ratio") with gr.Row(): output_zip = gr.File(label="Download Processed Images as ZIP") processing_time = gr.Textbox(label="Processing Time (seconds)") bg_choice.change(show_color_picker, inputs=bg_choice, outputs=custom_color) process_button.click(process, inputs=[input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers], outputs=[gallery_processed, output_zip, processing_time]) gallery_processed.select(update_compare, outputs=[image_original, image_processed, original_ratio, processed_ratio]) iface.launch(share=True)