Generative-AI / RemoveBG_By_SECourses.py
MonsterMMORPG's picture
Upload RemoveBG_By_SECourses.py
bb1de4c verified
raw
history blame
7.81 kB
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 V2 '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("Single Image Process")
batch_button = gr.Button("Batch Process Images in Given Folder")
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)