Upload RemoveBG_By_SECourses.py

RemoveBG_By_SECourses.py

@@ -0,0 +1,200 @@
+import os
+import cv2
+import numpy as np
+import torch
+import gradio as gr
+import argparse
+from pathlib import Path
+from glob import glob
+from typing import Optional, Tuple, List
+from PIL import Image
+from transformers import AutoModelForImageSegmentation
+from torchvision import transforms
+import time
+import os
+import platform
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Run the image segmentation app")
+    parser.add_argument("--share", action="store_true", help="Enable sharing of the Gradio interface")
+    return parser.parse_args()
+
+torch.set_float32_matmul_precision('high')
+torch.jit.script = lambda f: f
+
+os.environ['HOME'] = os.path.expanduser('~')
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def open_folder():
+    open_folder_path = os.path.abspath("results")
+    if platform.system() == "Windows":
+        os.startfile(open_folder_path)
+    elif platform.system() == "Linux":
+        os.system(f'xdg-open "{open_folder_path}"')
+
+class ImagePreprocessor():
+    def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
+        self.transform_image = transforms.Compose([
+            transforms.ToTensor(),
+        ])
+        self.normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+
+    def proc(self, image: Image.Image) -> torch.Tensor:
+        image = image.convert('RGB')  # Convert to RGB
+        image = self.transform_image(image)
+        return self.normalize(image)
+
+usage_to_weights_file = {
+    'General': 'BiRefNet',
+    'General-Lite': 'BiRefNet_T',
+    'Portrait': 'BiRefNet-portrait',
+    'DIS': 'BiRefNet-DIS5K',
+    'HRSOD': 'BiRefNet-HRSOD',
+    'COD': 'BiRefNet-COD',
+    'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs'
+}
+
+birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True)
+birefnet.to(device)
+birefnet.eval()
+
+def process_single_image(image_path: str, resolution: str, output_folder: str) -> Tuple[str, str, float]:
+    start_time = time.time()
+    
+    image = Image.open(image_path).convert('RGBA')
+    
+    if resolution == '':
+        resolution = f"{image.width}x{image.height}"
+    resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
+    
+    image_shape = image.size[::-1]
+    image_pil = image.resize(tuple(resolution))
+
+    image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
+    image_proc = image_preprocessor.proc(image_pil)
+    image_proc = image_proc.unsqueeze(0)
+
+    with torch.no_grad():
+        scaled_pred_tensor = birefnet(image_proc.to(device))[-1].sigmoid()
+
+    if device == 'cuda':
+        scaled_pred_tensor = scaled_pred_tensor.cpu()
+    
+    pred = torch.nn.functional.interpolate(scaled_pred_tensor, size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy()
+
+    pred_rgba = np.zeros((*pred.shape, 4), dtype=np.uint8)
+    pred_rgba[..., :3] = (pred[..., np.newaxis] * 255).astype(np.uint8)
+    pred_rgba[..., 3] = (pred * 255).astype(np.uint8)
+
+    image_array = np.array(image)
+    image_pred = image_array * (pred_rgba / 255.0)
+    
+    output_image = Image.fromarray(image_pred.astype(np.uint8), 'RGBA')
+    
+    base_filename = os.path.splitext(os.path.basename(image_path))[0]
+    output_path = os.path.join(output_folder, f"{base_filename}.png")
+    
+    counter = 1
+    while os.path.exists(output_path):
+        output_path = os.path.join(output_folder, f"{base_filename}_{counter:04d}.png")
+        counter += 1
+
+    output_image.save(output_path)
+    
+    processing_time = time.time() - start_time
+    print(f"Processed {image_path} in {processing_time:.4f} seconds")  # Added this line to print processing time
+    return image_path, output_path, processing_time
+
+def predict(
+    image: str,
+    resolution: str,
+    weights_file: Optional[str],
+    batch_folder: Optional[str] = None,
+    output_folder: Optional[str] = None,
+    is_batch: bool = False
+) -> Tuple[str, List[Tuple[str, str]]]:
+    global birefnet
+    _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
+    print('Using weights:', _weights_file)
+    birefnet = AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True)
+    birefnet.to(device)
+    birefnet.eval()
+
+    if not output_folder:
+        output_folder = 'results'
+    os.makedirs(output_folder, exist_ok=True)
+
+    results = []
+
+    if is_batch and batch_folder:
+        image_files = glob(os.path.join(batch_folder, '*'))
+        total_images = len(image_files)
+        processed_images = 0
+        start_time = time.time()
+
+        for img_path in image_files:
+            try:
+                input_path, output_path, proc_time = process_single_image(img_path, resolution, output_folder)
+                results.append((output_path, f"{proc_time:.4f} seconds"))
+                processed_images += 1
+                elapsed_time = time.time() - start_time
+                avg_time_per_image = elapsed_time / processed_images
+                estimated_time_left = avg_time_per_image * (total_images - processed_images)
+
+                status = f"Processed {processed_images}/{total_images} images. Estimated time left: {estimated_time_left:.2f} seconds"
+                print(status)
+            except Exception as e:
+                print(f"Error processing {img_path}: {str(e)}")
+                continue
+
+        return f"Batch processing complete. Processed {processed_images}/{total_images} images.", results
+    else:
+        input_path, output_path, proc_time = process_single_image(image, resolution, output_folder)
+        results.append((output_path, f"{proc_time:.4f} seconds"))
+        return "Single image processing complete.", results
+
+def create_interface():
+    with gr.Blocks() as demo:
+        gr.Markdown("## SECourses Improved BiRefNet V1 'Bilateral Reference for High-Resolution Dichotomous Image Segmentation' APP - SOTA Background Remover")
+        gr.Markdown("## Most Advanced Latest Version On : https://www.patreon.com/posts/109913645")
+        
+        with gr.Row():
+            input_image = gr.Image(type="filepath", label="Input Image",height=512)
+            output_image = gr.Gallery(label="Output Image", elem_id="gallery",height=512)
+
+
+        with gr.Row():
+            resolution = gr.Textbox(label="Resolution", placeholder="1024x1024 - Optional - Don't enter to use original image resolution - Higher res uses more VRAM but still works perfect with shared VRAM so fast")
+            weights_file = gr.Dropdown(choices=list(usage_to_weights_file.keys()), value="General", label="Weights File")
+            btn_open_outputs = gr.Button("Open Results Folder")
+            btn_open_outputs.click(fn=open_folder)
+
+        with gr.Row():
+            batch_folder = gr.Textbox(label="Batch Folder Path")
+            output_folder = gr.Textbox(label="Output Folder Path", value="results")
+
+        with gr.Row():
+            submit_button = gr.Button("Process")
+            batch_button = gr.Button("Process Batch")
+
+        output_text = gr.Textbox(label="Processing Status")
+
+        submit_button.click(
+            predict,
+            inputs=[input_image, resolution, weights_file, batch_folder, output_folder, gr.Checkbox(value=False, visible=False)],
+            outputs=[output_text, output_image]
+        )
+
+        batch_button.click(
+            predict,
+            inputs=[input_image, resolution, weights_file, batch_folder, output_folder, gr.Checkbox(value=True, visible=False)],
+            outputs=[output_text, output_image]
+        )
+
+    return demo
+
+if __name__ == "__main__":
+    args = parse_args()
+    demo = create_interface()
+    demo.launch(inbrowser=True, share=args.share)