Update app.py

app.py

@@ -1,94 +1,105 @@
-import torch
 import numpy as np
+import torch
 import torch.nn.functional as F
 from torchvision.transforms.functional import normalize
 from huggingface_hub import hf_hub_download
 import gradio as gr
-from PIL import Image
+from gradio_imageslider import ImageSlider
 from briarmbg import BriaRMBG
+import PIL
+from PIL import Image
 from typing import Tuple
 
-# Load the BriaRMBG model
-net = BriaRMBG()
+net=BriaRMBG()
+# model_path = "./model1.pth"
 model_path = hf_hub_download("briaai/RMBG-1.4", 'model.pth')
 if torch.cuda.is_available():
     net.load_state_dict(torch.load(model_path))
-    net = net.cuda()
+    net=net.cuda()
 else:
-    net.load_state_dict(torch.load(model_path, map_location="cpu"))
-net.eval()
-
-def resize_image(image) -> Image.Image:
-    """Resize image to fit model's input requirements."""
-    pixel_number = 960 * 960
-    granularity_val = 64
-    ratio = image.size[0] / image.size[1]
-    width = int((pixel_number * ratio) ** 0.5)
-    width -= width % granularity_val
-    height = int(pixel_number / width)
-    height -= height % granularity_val
-    return image.resize((width, height))
-
-def get_masked_background_image(image, image_mask) -> Tuple[np.ndarray, np.ndarray]:
-    """Apply the segmentation mask to the original image."""
-    image_mask = image_mask.resize(image.size) 
-    image = np.array(image.convert("RGB")).transpose(2, 0, 1).astype(np.float32) / 255.0
-    image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0
-    image[:, image_mask < 0.5] = 0  
-    return image, image_mask
-
-def get_control_image_tensor(vae, image, mask) -> torch.Tensor:
-    """Prepare the masked image tensor for model input."""
-    masked_image, image_mask = get_masked_background_image(image, mask)
-    masked_image_tensor = torch.from_numpy(masked_image)
-    masked_image_tensor = (masked_image_tensor - 0.5) / 0.5
-    masked_image_tensor = masked_image_tensor.unsqueeze(0).to(device="cuda:0")
-    control_latents = vae.encode(masked_image_tensor[:, :3, :, :].to(vae.dtype)).latent_dist.sample()
-    control_latents = control_latents * vae.config.scaling_factor
-    mask_tensor = torch.tensor(image_mask, dtype=torch.float32)[None, None, ...].to(device="cuda:0")
-    mask_resized = torch.nn.functional.interpolate(mask_tensor, size=(control_latents.shape[2], control_latents.shape[3]), mode='nearest')
-    control_tensor = torch.cat([control_latents, mask_resized], dim=1)
-    return control_tensor
+    net.load_state_dict(torch.load(model_path,map_location="cpu"))
+net.eval() 
 
-def remove_bg_from_image(image) -> Image.Image:
-    """Use BriaRMBG to generate a segmentation mask."""
-    from transformers import pipeline
-    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
-    mask = pipe(image, return_mask=True)
-    return mask
+    
+def resize_image(image):
+    image = image.convert('RGB')
+    model_input_size = (1024, 1024)
+    image = image.resize(model_input_size, Image.BILINEAR)
+    return image
 
-def paste_fg_over_image(gen_image: Image.Image, orig_image: Image.Image, fg_mask: Image.Image) -> Image.Image:
-    """Paste the foreground over the generated image."""
-    fg_mask = fg_mask.convert("L").resize(orig_image.size, Image.NEAREST)
-    gen_image = gen_image.convert("RGBA")
-    orig_image = orig_image.convert("RGBA")
-    gen_image.paste(orig_image, (0, 0), fg_mask)
-    return gen_image.convert("RGB")
 
 def process(image):
-    """Process image for background removal and pasting over new background."""
+
+    # prepare input
     orig_image = Image.fromarray(image)
+    w,h = orig_im_size = orig_image.size
     image = resize_image(orig_image)
-    mask = remove_bg_from_image(image)
-    result_image = paste_fg_over_image(image, orig_image, mask)
-    return result_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
+    im_array = (result*255).cpu().data.numpy().astype(np.uint8)
+    pil_im = Image.fromarray(np.squeeze(im_array))
+    # paste the mask on the original image
+    new_im = Image.new("RGBA", pil_im.size, (0,0,0,0))
+    new_im.paste(orig_image, mask=pil_im)
+    # new_orig_image = orig_image.convert('RGBA')
+
+    return new_im
+    # return [new_orig_image, new_im]]
+    
+# block = gr.Blocks().queue()
+
+# with block:
+#     gr.Markdown("## HBS_V1")
+#     gr.HTML('''
+        # <p style="margin-bottom: 10px; font-size: 94%">
+        #     This is a demo for Human Body Segmentation that using
+        #     YoloV8 image instance model as backbone. 
+        #   </p>
+#     ''')
+#     with gr.Row():
+#         with gr.Column():
+#             input_image = gr.Image(sources=None, type="pil") # None for upload, ctrl+v and webcam
+#             # input_image = gr.Image(sources=None, type="numpy") # None for upload, ctrl+v and webcam
+#             run_button = gr.Button(value="Run")
+            
+#         with gr.Column():
+#             result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", columns=[1], height='auto')
+#     ips = [input_image]
+#     run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
 
-# Set up Gradio interface
+# block.launch(debug = True)
+
+# block = gr.Blocks().queue()
+
+gr.Markdown("## HBS_V1")
+gr.HTML('''
+  <p style="margin-bottom: 10px; font-size: 94%">
+    This is a demo for Human Body Segmentation that using
+    YoloV8 image instance model as backbone. 
+  </p>
+''')
 title = "Human Body Segmentation"
-description = """
-Human Body Segmentation model by <a href='https://github.com/WildanJR09' target='_blank'><b>WildanJR</b></a>.
-Separate foreground and background for various image categories. Trained for commercial content.
+description = r"""Human Body Segmentation model developed by <a href='https://github.com/WildanJR09' target='_blank'><b>WildanJR</b></a>, Designed to effectively separate foreground from background in a range of categories and image types.<br> 
+This model has been trained on a carefully selected dataset, which includes: general stock images, e-commerce, gaming, and advertising content, making it suitable for commercial use cases powering enterprise content creation at scale. The accuracy, efficiency, and versatility currently rival leading source-available models. It is ideal where content safety, legally licensed datasets, and bias mitigation are paramount. For test upload your image and wait. </a>.<br>
 """
-examples = [['./jisoo.jpg']]
-
-demo = gr.Interface(
-    fn=process,
-    inputs="image",
-    outputs="image",
-    title=title,
-    description=description,
-    examples=examples
-)
+examples = [['./jisoo.jpg'],]
+# output = ImageSlider(position=0.5,label='Image without background', type="pil", show_download_button=True)
+# demo = gr.Interface(fn=process,inputs="image", outputs=output, examples=examples, title=title, description=description)
+demo = gr.Interface(fn=process,inputs="image", outputs="image", examples=examples, title=title, description=description)
 
 if __name__ == "__main__":
-    demo.launch(share=False)
+    demo.launch(share=False)