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)