import os, re import zipfile import shutil import time from PIL import Image, ImageDraw, ImageFont 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 class LoadModel: def __init__(self): self.device = "cuda" if torch.cuda.is_available() else "cpu" def get_stable_diffusion_model(self): return StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float32).to(self.device) def get_bria_model(self): return pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True, device=self.device) def remove_background_bria(self, input_path, pipeline): print(f"Removing background using bria for image: {input_path}") result = pipeline(input_path) return result def remove_background_rembg(self, 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 class ImageProcessor: def __init__(self): self.model_loader = LoadModel() self.bria_pipeline = self.model_loader.get_bria_model() def get_bounding_box_with_threshold(self, 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 top, bottom = np.where(rows)[0][[0, -1]] left, right = np.where(cols)[0][[0, -1]] if left < right and top < bottom: return (left, top, right, bottom) else: return None def position_logic(self, image_path, canvas_size, padding_top, padding_right, padding_bottom, padding_left, use_threshold=True): print('masuk ke method position_logic') image = Image.open(image_path) image = image.convert("RGBA") print('image converted to RGBA') # Get the bounding box of the non-blank area with threshold if use_threshold: bbox = self.get_bounding_box_with_threshold(image, threshold=10) else: bbox = image.getbbox() log = [] if bbox: # Check 1 pixel around the image for non-transparent pixels width, height = image.size cropped_sides = [] # Define tolerance for transparency tolerance = 30 # Adjust this value as needed # Check top edge if any(image.getpixel((x, 0))[3] > tolerance for x in range(width)): cropped_sides.append("top") # Check bottom edge if any(image.getpixel((x, height - 1))[3] > tolerance for x in range(width)): cropped_sides.append("bottom") # Check left edge if any(image.getpixel((0, y))[3] > tolerance for y in range(height)): cropped_sides.append("left") # Check right edge if any(image.getpixel((width - 1, y))[3] > tolerance for y in range(height)): cropped_sides.append("right") if cropped_sides: info_message = f"Info for {os.path.basename(image_path)}: The following sides of the image may contain cropped objects: {', '.join(cropped_sides)}" print(info_message) log.append({"info": info_message}) else: info_message = f"Info for {os.path.basename(image_path)}: The image is not cropped." print(info_message) print("ini nih cropped side",cropped_sides) 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 the new size to expand the image target_width, target_height = canvas_size aspect_ratio = image.width / image.height if len(cropped_sides) == 4: # If the image is cropped on all sides, center crop it to fit the canvas 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) image = image.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) image = image.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 print(cropped_sides) elif not cropped_sides: # If the image is not cropped, expand it from center until it touches the padding new_height = target_height - padding_top - padding_bottom new_width = int(new_height * aspect_ratio) if new_width > target_width - padding_left - padding_right: # If width exceeds available space, adjust based on width new_width = target_width - padding_left - padding_right new_height = int(new_width / aspect_ratio) # Resize the image 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: # Logic for handling cropped images new_height = target_height - padding_bottom new_width = int(new_height * aspect_ratio) # If new width exceeds canvas width, adjust based on width if new_width > target_width: new_width = target_width new_height = int(new_width / aspect_ratio) # Resize the image image = image.resize((new_width, new_height), Image.LANCZOS) log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)}) # Set position if "left" in cropped_sides: x = 0 else: x = target_width - new_width y = 0 return log, image, x, y def process_single_image(self, image_path, output_folder, bg_method, canvas_size_name, output_format, bg_choice, custom_color, watermark_path=None): print('masuk ke method process_single_image') add_padding_line = False if canvas_size_name == 'Rox': canvas_size = (1080, 1080) padding_top = 112 padding_right = 125 padding_bottom = 116 padding_left = 125 elif canvas_size_name == 'Columbia': canvas_size = (730, 610) padding_top = 30 padding_right = 105 padding_bottom = 35 padding_left = 105 elif canvas_size_name == 'Zalora': canvas_size = (763, 1100) padding_top = 50 padding_right = 50 padding_bottom = 200 padding_left = 50 filename = os.path.basename(image_path) try: print(f"Processing image: {filename}") if bg_method == 'rembg': image_with_no_bg = self.model_loader.remove_background_rembg(image_path) elif bg_method == 'bria': image_with_no_bg = self.model_loader.remove_background_bria(image_path, self.bria_pipeline) elif bg_method == None: image_with_no_bg = Image.open(image_path) 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 = self.position_logic(temp_image_path, canvas_size, padding_top, padding_right, padding_bottom, padding_left) # Create a new canvas with the appropriate background 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 visible black line for padding when background is not transparent if add_padding_line: draw = ImageDraw.Draw(canvas) draw.rectangle([padding_left, padding_top, canvas_size[0] - padding_right, canvas_size[1] - padding_bottom], outline="black", width=5) log.append({"action": "add_padding_line"}) 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) # Apply watermark only if the filename ends with "_01" and watermark_path is provided if os.path.splitext(filename)[0].endswith("_01") and watermark_path: watermark = Image.open(watermark_path).convert("RGBA") canvas = canvas.convert("RGBA") canvas.paste(watermark, (0, 0), watermark) log.append({"action": "add_watermark"}) if output_format == 'JPG': canvas = canvas.convert('RGB') canvas.save(output_path, format='JPEG') else: canvas.save(output_path, format='PNG') os.remove(temp_image_path) print(f"Processed image path: {output_path}") return [(output_path, image_path)], log except Exception as e: print(f"Error processing {filename}: {e}") return None, None def remove_extension(self, filename): # Regular expression to match any extension at the end of the string return re.sub(r'\.[^.]+$', '', filename) def process_images(self, 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()): print('masuk ke method process_images') 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 = [] 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) except zipfile.BadZipFile as e: print(f"Error extracting zip file: {e}") return [], None, 0 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'))] 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) print(f"Total images to process: {total_images}") avg_processing_time = 0 with ThreadPoolExecutor(max_workers=num_workers) as executor: future_to_image = {executor.submit(self.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: if watermark_path: get_name = future_to_image[future].split('/') get_name = self.remove_extension(get_name[len(get_name)-1]) twibbon_input = f'{get_name}.png' if output_format == 'PNG' else f'{get_name}.jpg' twibbon_output_path = os.path.join(output_folder, f'result_{start_time_image}.png') print('mencoba method add twibon') self.add_twibbon(f'processed_images/{twibbon_input}', watermark_path, twibbon_output_path) print('method add_twibon berhasil') processed_images.append((twibbon_output_path, twibbon_output_path)) else: 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: print('HAYOLOH KETAUAN ERRORNYA') print(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) print("Comprehensive log saved to", os.path.join(output_folder, 'process_log.json')) end_time = time.time() processing_time = end_time - start_time print(f"Processing time: {processing_time} seconds") return original_images, processed_images, output_zip_path, processing_time def remove_white_background(self, twibbon, tolerance=100): """ Menghapus background putih dengan toleransi tertentu. tolerance: Nilai antara 0 (tidak toleran, hanya putih murni) hingga 255 (sangat toleran, mencakup hampir semua warna cerah). """ twibbon = twibbon.convert("RGBA") data = twibbon.getdata() new_data = [] for item in data: # Hitung jarak warna ke putih (255, 255, 255) distance_to_white = sum([abs(255 - c) for c in item[:3]]) # RGB distance if distance_to_white <= tolerance: # Jika jarak warna ke putih lebih kecil dari toleransi, buat transparan new_data.append((255, 255, 255, 0)) # Transparan penuh else: # Tetap pertahankan warna asli new_data.append(item) twibbon.putdata(new_data) return twibbon def adjust_opacity(self, twibbon, opacity_level): twibbon = twibbon.convert("RGBA") data = twibbon.getdata() new_data = [] for item in data: # Ubah hanya nilai alpha (transparansi) new_alpha = int(item[3] * opacity_level / 255) # Sesuaikan alpha sesuai opacity_level new_data.append((item[0], item[1], item[2], new_alpha)) twibbon.putdata(new_data) return twibbon def add_twibbon(self, image_path, twibbon_path, output_path): # Open the original image image = Image.open(image_path).convert("RGBA") print('Original image loaded') # Open the twibbon (watermark) twibbon = Image.open(twibbon_path).convert("RGBA") print('Twibbon (watermark) loaded') # Remove white background from twibbon twibbon = self.remove_white_background(twibbon) # twibbon = self.adjust_opacity(twibbon, 128) # Resize the twibbon (watermark) image_width, image_height = image.size twibbon_size = (image_width // 5, image_height // 5) # Resize twibbon to 20% of image size twibbon = twibbon.resize(twibbon_size, Image.Resampling.LANCZOS) # Center the watermark twibbon_width, twibbon_height = twibbon.size x_offset = (image_width - twibbon_width) // 2 y_offset = (image_height - twibbon_height) // 2 # Create a new transparent layer for the watermark transparent_layer = Image.new("RGBA", (image_width, image_height), (0, 0, 0, 0)) transparent_layer.paste(twibbon, (x_offset, y_offset), mask=twibbon.split()[3]) # Composite the image with the transparent layer final_image = Image.alpha_composite(image, transparent_layer) # Save the final result print('Saving the final image with watermark...') final_image.save(output_path) print('Image saved successfully') return final_image class ModelInference: def __init__(self): self.loader = LoadModel() self.sd_model = self.loader.get_stable_diffusion_model() def text_to_image(self, prompt): os.makedirs("generated_images", exist_ok=True) # Ensure the directory exists image = self.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(self, 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 = self.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 class CreativeImageSuite: def __init__(self): self.inference = ModelInference() # Use the ModelInference class self.processor = ImageProcessor() self.theme = "NoCrypt/miku@1.2.2" self.title = "# 🎨 Creative Image Suite: Generate, Modify, and Enhance Your Visuals" self.description = """ **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. 2. **🖼️ Image to Image**: Enhance existing images by providing a text description of the modifications you want. 3. **🖌️ Image Background Removal and Resizing**: Effortlessly remove backgrounds from images, resize them, and even add watermarks (optional). """ self.interface = None def gradio_interface(self, input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers): progress = gr.Progress() watermark_path = watermark.name if watermark else None # Check input_files, is it single image, list image, or zip/rar if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')): return self.processor.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 self.processor.process_images(input_files, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress) else: return self.processor.process_images(input_files.name, bg_method, watermark_path, canvas_size, output_format, bg_choice, custom_color, num_workers, progress) def show_color_picker(self, bg_choice): if bg_choice == 'custom': return gr.update(visible=True) return gr.update(visible=False) def update_compare(self, evt: gr.SelectData): if isinstance(evt.value, dict) and 'caption' in evt.value: input_path = evt.value['caption'] output_path = evt.value['image']['path'] input_path = input_path.split("Input: ")[-1] # 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: print("No caption found in selection") return gr.update(value=None), gr.update(value=None), gr.update(value=None), gr.update(value=None) def master_process(self, input_files, bg_method, watermark, canvas_size, output_format, bg_choice, custom_color, num_workers): _, processed_images, zip_path, time_taken = self.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" def build_interface(self): with gr.Blocks(theme=self.theme) as self.interface: # App Header gr.Markdown(self.title) gr.Markdown(self.description) # Text-to-Image Section gr.Markdown("## Text to Image Feature") gr.Markdown(""" *Create Visuals from Your Imagination* This feature allows you to generate unique images from a simple text description. Just type your idea in the box, and our tool will bring it to life in seconds. - Example Prompts: - "Generate an image of a kitchenset with modern furniture." - "Generate an image of a casual turtleneck shirt." - "A futuristic laptop on a white background ." """) with gr.Row(): prompt_input = gr.Textbox(label="Describe your vision to generate a unique image:") generate_button = gr.Button("Create Image") output_image = gr.Image(label="Your Generated Image") download_button = gr.File(label="Download Image", type="filepath") generate_button.click(self.inference.text_to_image, inputs=prompt_input, outputs=[output_image, download_button]) # Image-to-Image Section gr.Markdown("## Image to Image Feature") gr.Markdown(""" *Enhance or Transform Your Existing Images* Upload an image and describe the changes you'd like to see. From subtle edits to artistic transformations, this feature helps you bring new life to your visuals. - Example Edits: - "Change the Color of the bag" - "Add plastic to the suitcase." - "mirror this photo of denim jacket." """) with gr.Row(): input_image = gr.Image(label="Upload an image to modify:", type="pil") prompt_modification = gr.Textbox(label="Describe the changes you want:") modify_button = gr.Button("Apply Changes") modified_output_image = gr.Image(label="Your Modified Image") download_modified_button = gr.File(label="Download Modified Image", type="filepath") modify_button.click(self.inference.text_image_to_image, inputs=[input_image, prompt_modification], outputs=[modified_output_image, download_modified_button]) # Background Removal and Resizing Section gr.Markdown("## Image Background Removal and Resizing with Optional Watermark") gr.Markdown(""" *Perfect Your Images for Any Use Case* Easily remove backgrounds, resize images to fit your needs, and even add watermarks to maintain originality or branding. This feature is ideal for e-commerce, social media, and design projects. - Features: - Supports batch processing of multiple images or ZIP/RAR files. - Options for transparent, solid color, or custom backgrounds. - Output in your choice of PNG or JPG format. """) with gr.Row(): input_files = gr.File(label="Upload an image or a ZIP/RAR file for batch processing:", file_types=[".zip", ".rar", "image"], interactive=True) watermark = gr.File(label="Upload an optional watermark (PNG only):", file_types=[".png"]) with gr.Row(): canvas_size = gr.Radio(choices=["Rox", "Columbia", "Zalora"], label="Select the desired canvas size:", value="Rox") output_format = gr.Radio(choices=["PNG", "JPG"], label="Choose the output format:", value="JPG") num_workers = gr.Slider(minimum=1, maximum=16, step=1, label="Set the number of processing threads:", value=5) with gr.Row(): bg_method = gr.Radio(choices=["bria", "rembg", None], label="Choose a background removal method:", value="bria") bg_choice = gr.Radio(choices=["transparent", "white", "custom"], label="Select a background style:", value="white") custom_color = gr.ColorPicker(label="Pick a custom background color (if applicable):", value="#ffffff", visible=False) process_button = gr.Button("Start Processing") with gr.Row(): gallery_processed = gr.Gallery(label="Processed Images") with gr.Row(): image_original = gr.Image(label="Original Image Preview", interactive=False) image_processed = gr.Image(label="Processed Image Preview", interactive=False) with gr.Row(): original_ratio = gr.Textbox(label="Aspect Ratio (Original)") processed_ratio = gr.Textbox(label="Aspect Ratio (Processed)") with gr.Row(): output_zip = gr.File(label="Download All Processed Images (ZIP)") processing_time = gr.Textbox(label="Total Processing Time (seconds)") bg_choice.change(self.show_color_picker, inputs=bg_choice, outputs=custom_color) process_button.click(self.master_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(self.update_compare, outputs=[image_original, image_processed, original_ratio, processed_ratio]) def launch(self): if self.interface is None: self.build_interface() self.interface.launch(share=True, debug=True) app = CreativeImageSuite() app.launch()