Update app.py

app.py

@@ -1,64 +1,108 @@
-import gradio as gr
-import threading
-import os
+import numpy as np
 import torch
+import torch.nn.functional as F
+from torchvision.transforms.functional import normalize
+import gradio as gr
+from gradio_imageslider import ImageSlider
+from briarmbg import BriaRMBG
+import PIL
+from PIL import Image
+from typing import Tuple
+import cv2
+
+# Load the model
+net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+net.to(device)
+net.eval()
+
+def resize_image(image):
+    image = image.convert('RGB')
+    model_input_size = (1024, 1024)
+    image = image.resize(model_input_size, Image.BILINEAR)
+    return image
+
+def process_image(image):
+    # prepare input
+    orig_image = Image.fromarray(image)
+    w, h = orig_im_size = orig_image.size
+    image = resize_image(orig_image)
+    im_np = np.array(image)
+    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1)
+    im_tensor = torch.unsqueeze(im_tensor, 0)
+    im_tensor = torch.divide(im_tensor, 255.0)
+    im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
+    if torch.cuda.is_available():
+        im_tensor = im_tensor.cuda()
+
+    # inference
+    result = net(im_tensor)
+    # post process
+    result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear'), 0)
+    ma = torch.max(result)
+    mi = torch.min(result)
+    result = (result - mi) / (ma - mi)
+    # image to pil
+    result_array = (result * 255).cpu().data.numpy().astype(np.uint8)
+    pil_mask = Image.fromarray(np.squeeze(result_array))
+    # add the mask on the original image as alpha channel
+    new_im = orig_image.copy()
+    new_im.putalpha(pil_mask)
+    return new_im
 
-os.environ["OMP_NUM_THREADS"] = str(os.cpu_count())
-torch.set_num_threads(os.cpu_count())
+def process_video(video_path):
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise ValueError("Error opening video file")
 
-model1 = gr.load("models/prithivMLmods/SD3.5-Turbo-Realism-2.0-LoRA")
-model2 = gr.load("models/Purz/face-projection")
+    # Get video properties
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter('output.mp4', fourcc, fps, (width, height), isColor=True)
 
-stop_event = threading.Event()
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
 
-def generate_images(text, selected_model):
-    stop_event.clear()
+        # Convert frame to PIL Image
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        pil_image = Image.fromarray(frame)
 
-    if selected_model == "Model 1 (Turbo Realism)":
-        model = model1
-    elif selected_model == "Model 2 (Face Projection)":
-        model = model2
-    else:
-        return ["Invalid model selection."] * 3
+        # Process the frame
+        processed_image = process_image(frame)
 
-    results = []
-    for i in range(3):
-        if stop_event.is_set():
-            return ["Image generation stopped by user."] * 3
+        # Convert back to OpenCV format
+        processed_frame = cv2.cvtColor(np.array(processed_image), cv2.COLOR_RGBA2BGRA)
 
-        modified_text = f"{text} variation {i+1}"
-        result = model(modified_text)
-        results.append(result)
+        # Write the frame to the output video
+        out.write(processed_frame)
 
-    return results
+    cap.release()
+    out.release()
+    return 'output.mp4'
 
-def stop_generation():
-    """Stops the ongoing image generation by setting the stop_event flag."""
-    stop_event.set()
-    return ["Generation stopped."] * 3
+def process_input(input_data):
+    if isinstance(input_data, str):  # Assuming video path is provided as a string
+        return process_video(input_data)
+    else:  # Assuming image is provided as numpy array
+        return process_image(input_data)
 
-with gr.Blocks() as interface:#...
-    gr.Markdown(
-        "### ⚠ Sorry for the inconvenience. The Space is currently running on the CPU, which might affect performance. We appreciate your understanding."
-    )
-    
-    text_input = gr.Textbox(label="Type here your imagination:", placeholder="Type your prompt...")
-    model_selector = gr.Radio(
-        ["Model 1 (Turbo Realism)", "Model 2 (Face Projection)"],
-        label="Select Model",
-        value="Model 1 (Turbo Realism)"
-    )
-    
-    with gr.Row():
-        generate_button = gr.Button("Generate 3 Images 🎨")
-        stop_button = gr.Button("Stop Image Generation")
-    
-    with gr.Row():
-        output1 = gr.Image(label="Generated Image 1")
-        output2 = gr.Image(label="Generated Image 2")
-        output3 = gr.Image(label="Generated Image 3")
-    
-    generate_button.click(generate_images, inputs=[text_input, model_selector], outputs=[output1, output2, output3])
-    stop_button.click(stop_generation, inputs=[], outputs=[output1, output2, output3])
+gr.Markdown("## BRIA RMBG 1.4")
+gr.HTML('''
+  <p style="margin-bottom: 10px; font-size: 94%">
+    This is a demo for BRIA RMBG 1.4 that using
+    <a href="https://huggingface.co/briaai/RMBG-1.4" target="_blank">BRIA RMBG-1.4 image matting model</a> as backbone. 
+  </p>
+''')
+title = "Background Removal"
+description = r"""Background removal model developed by <a href='https://BRIA.AI' target='_blank'><b>BRIA.AI</b></a>, trained on a carefully selected dataset and is available as an open-source model for non-commercial use.<br> 
+For test upload your image and wait. Read more at model card <a href='https://huggingface.co/briaai/RMBG-1.4' target='_blank'><b>briaai/RMBG-1.4</b></a>. To purchase a commercial license, simply click <a href='https://go.bria.ai/3ZCBTLH' target='_blank'><b>Here</b></a>. <br>
+"""
+examples = [['./input.jpg'],]
+demo = gr.Interface(fn=process_input, inputs="file", outputs="playable_video", examples=examples, title=title, description=description)
 
-interface.launch()
+if __name__ == "__main__":
+    demo.launch(share=False)