Update app.py

app.py

@@ -32,7 +32,7 @@ def remove_background_rembg(input_path):
 
 def remove_background_bria(input_path):
     print(f"Removing background using bria for image: {input_path}")
-    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True, device=1)
+    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True, device=0)
     pillow_image = pipe(input_path)
     return pillow_image
 
@@ -52,58 +52,26 @@ def text_image_to_image(image, prompt):
     modified_image.save(image_path)
     return modified_image, image_path
 
-def process_images(image_paths, remove_bg_method):
-    with ThreadPoolExecutor() as executor:
-        if remove_bg_method == 'rembg':
-            results = list(executor.map(remove_background_rembg, image_paths))
-        elif remove_bg_method == 'bria':
-            results = list(executor.map(remove_background_bria, image_paths))
-        else:  # No background removal
-            results = image_paths  # Just return the original paths
-    return results
-
 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)
     
-    try:
-        top, bottom = np.where(rows)[0][[0, -1]]
-        left, right = np.where(cols)[0][[0, -1]]
-    except IndexError:
-        return None  # No non-transparent pixels found
+    # Find the bounding box
+    top, bottom = np.where(rows)[0][[0, -1]]
+    left, right = np.where(cols)[0][[0, -1]]
     
-    return (left, top, right, bottom) if left < right and top < bottom else None
-
-def check_cropped_sides(image, tolerance):
-    cropped_sides = []
-    width, height = image.size
-    edges = {
-        "top": [(x, 0) for x in range(width)],
-        "bottom": [(x, height - 1) for x in range(width)],
-        "left": [(0, y) for y in range(height)],
-        "right": [(width - 1, y) for y in range(height)]
-    }
-
-    for side, pixels in edges.items():
-        if any(image.getpixel(pixel)[3] > tolerance for pixel in pixels):
-            cropped_sides.append(side)
-
-    return cropped_sides
-
-def resize_image(image, target_size, aspect_ratio):
-    target_width, target_height = target_size
-    if aspect_ratio > 1:  # Landscape
-        new_height = target_height
-        new_width = int(new_height * aspect_ratio)
-    else:  # Portrait or square
-        new_width = target_width
-        new_height = int(new_width / aspect_ratio)
-
-    return image.resize((new_width, new_height), Image.LANCZOS), new_width, new_height
-
+    if left < right and top < bottom:
+        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):
     image = Image.open(image_path)
     image = image.convert("RGBA")
@@ -400,14 +368,6 @@ def position_logic(image_path, canvas_size, padding_top, padding_right, padding_
 
     return log, image, x, y
 
-def get_canvas_size(canvas_size_name):
-    sizes = {
-        'Rox': ((1080, 1080), (112, 125, 116, 125)),
-        'Columbia': ((730, 610), (30, 105, 35, 105)),
-        'Zalora': ((763, 1100), (50, 50, 200, 50)),
-    }
-    return sizes.get(canvas_size_name, ((1080, 1080), (0, 0, 0, 0)))
-
 def process_single_image(image_path, output_folder, bg_method, canvas_size_name, output_format, bg_choice, custom_color, watermark_path=None):
     add_padding_line = False
 
@@ -484,59 +444,6 @@ def process_single_image(image_path, output_folder, bg_method, canvas_size_name,
     except Exception as e:
         print(f"Error processing {filename}: {e}")
         return None, None
-
-def create_canvas(canvas_size, bg_choice, custom_color):
-    if bg_choice == 'white':
-        return Image.new("RGBA", canvas_size, "WHITE")
-    elif bg_choice == 'custom':
-        return Image.new("RGBA", canvas_size, custom_color)
-    else:  # transparent
-        return Image.new("RGBA", canvas_size, (0, 0, 0, 0))
-
-# Function to add watermark
-def generate_image(prompt):
-    # Code to generate an image using Stable Diffusion
-    # Placeholder: Replace with actual image generation code
-    response = requests.get("URL_OF_GENERATED_IMAGE")  # Replace with the actual URL or generation method
-    img = Image.open(BytesIO(response.content))
-    return img
-
-def add_watermark(image, watermark_path):
-    watermark = Image.open(watermark_path).convert("RGBA")
-    
-    # Resize watermark if needed
-    watermark = watermark.resize((image.width // 4, image.height // 4), Image.ANTIALIAS)
-    
-    # Position watermark in the center
-    position = ((image.width - watermark.width) // 2, (image.height - watermark.height) // 2)
-    
-    # Composite the watermark onto the image
-    transparent = Image.new('RGBA', image.size, (0, 0, 0, 0))
-    transparent.paste(image, (0, 0))
-    transparent.paste(watermark, position, watermark)
-    
-    return transparent
-
-def process_image(prompt, watermark_path, bg_removal_choice):
-    image = generate_image(prompt)
-    
-    if bg_removal_choice == "none":
-        final_image = image
-    else:
-        # Placeholder for actual background removal code
-        final_image = image  # Replace with background removal logic
-    
-    final_image_with_watermark = add_watermark(final_image, watermark_path)
-    return final_image_with_watermark
-        
-def save_image(canvas, output_folder, filename, output_format):
-    output_ext = 'jpg' if output_format == 'JPG' else 'png'
-    output_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}.{output_ext}")
-    if output_format == 'JPG':
-        canvas.convert('RGB').save(output_path, format='JPEG')
-    else:
-        canvas.save(output_path, format='PNG')
-    return output_path
     
 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()):
     start_time = time.time()
@@ -599,32 +506,7 @@ def process_images(input_files, bg_method='rembg', watermark_path=None, canvas_s
                 print(f"Error processing image {future_to_image[future]}: {e}")
 
     output_zip_path = "processed_images.zip"
-    # Handle uploaded files
-    if isinstance(input_files, list):
-        images = [Image.open(file.name) for file in input_files]
-    else:
-        with zipfile.ZipFile(input_files.name, 'r') as zip_ref:
-            zip_ref.extractall("extracted_images")
-            images = [Image.open(os.path.join("extracted_images", f)) for f in os.listdir("extracted_images")]
-
-    # Background removal logic (placeholder)
-    for img in images:
-        # Implement background removal based on selected method
-        if bg_method == "rembg":
-            img = remove_background_with_rembg(img)  # Placeholder function
-        elif bg_method == "bria":
-            img = remove_background_with_bria(img)  # Placeholder function
-
-        # Add watermark if provided
-        if watermark:
-            watermark_img = Image.open(watermark.name).convert("RGBA")
-            img = add_watermark(img, watermark_img)
-
-        # Save processed image
-        output_path = f"processed_{len(processed_images)}.{output_format.lower()}"
-        img.save(output_path, format=output_format)
-        processed_images.append(output_path)
-
+    
     # Create a zip file of processed images
     with zipfile.ZipFile(output_zip_path, 'w') as zipf:
         for processed_img in processed_images:
@@ -632,75 +514,6 @@ def process_images(input_files, bg_method='rembg', watermark_path=None, canvas_s
     
     return processed_images, output_zip_path
 
-def add_watermark(image, watermark):
-    # Resize and position watermark
-    watermark = watermark.resize((image.width // 4, image.height // 4), Image.ANTIALIAS)
-    position = ((image.width - watermark.width) // 2, (image.height - watermark.height) // 2)
-    
-    transparent = Image.new('RGBA', image.size, (0, 0, 0, 0))
-    transparent.paste(image.convert("RGBA"), (0, 0))
-    transparent.paste(watermark, position, watermark)
-    
-    return transparent.convert("RGB")
-
-def remove_background_with_rembg(image):
-    # Placeholder for actual background removal code
-    return image  # Replace with actual removal logic
-
-def remove_background_with_bria(image):
-    # Placeholder for actual background removal code
-    return image  # Replace with actual removal logic
-
-def show_color_picker(bg_choice):
-    return bg_choice == "custom"
-
-def update_compare(selected_image):
-    # Logic to display original vs processed images
-    return selected_image, selected_image  # Replace with actual images
-    
-    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 extract_image_files(input_files):
-    if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')):
-        input_folder = "temp_input"
-        if os.path.exists(input_folder):
-            shutil.rmtree(input_folder)
-        os.makedirs(input_folder)
-
-        with zipfile.ZipFile(input_files, 'r') as zip_ref:
-            zip_ref.extractall(input_folder)
-
-        return [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):
-        return input_files
-    else:
-        return [input_files]
-
-def create_output_zip(processed_images):
-    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))
-    return output_zip_path
-
-def save_log(all_logs, output_folder):
-    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'))
-
 def gradio_interface(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