Add BiRefNet-Lite-2K model, add model choosing option in tab_batch.

app.py

@@ -61,7 +61,7 @@ def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
 class ImagePreprocessor():
  def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
  self.transform_image = transforms.Compose([
- transforms.Resize(resolution), # 1. keep consistent with the cv2.resize used in training 2. redundant with that in path_to_image()
+ transforms.Resize(resolution),
  transforms.ToTensor(),
  transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
  ])
@@ -74,6 +74,7 @@ class ImagePreprocessor():
 usage_to_weights_file = {
  'General': 'BiRefNet',
  'General-Lite': 'BiRefNet_lite',
+ 'General-Lite-2K': 'BiRefNet_lite-2K',
  'Portrait': 'BiRefNet-portrait',
  'DIS': 'BiRefNet-DIS5K',
  'HRSOD': 'BiRefNet-HRSOD',
@@ -102,8 +103,8 @@ def predict(images, resolution, weights_file):
  try:
  resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
  except:
- resolution = [1024, 1024]
- print('Invalid resolution input. Automatically changed to 1024x1024.')
+ resolution = (1024, 1024) if weights_file not in ['General-Lite-2K'] else (2560, 1440)
+ print('Invalid resolution input. Automatically changed to 1024x1024 or 2K.')
 
  if isinstance(images, list):
  # For tab_batch
@@ -182,7 +183,7 @@ tab_image = gr.Interface(
  fn=predict,
  inputs=[
  gr.Image(label='Upload an image'),
- gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`. Higher resolutions can be much slower for inference.", label="Resolution"),
+ gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`.", label="Resolution"),
  gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.")
  ],
  outputs=ImageSlider(label="BiRefNet's prediction", type="pil"),
@@ -195,7 +196,7 @@ tab_text = gr.Interface(
  fn=predict,
  inputs=[
  gr.Textbox(label="Paste an image URL"),
- gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`. Higher resolutions can be much slower for inference.", label="Resolution"),
+ gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`.", label="Resolution"),
  gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.")
  ],
  outputs=ImageSlider(label="BiRefNet's prediction", type="pil"),
@@ -206,7 +207,11 @@ tab_text = gr.Interface(
 
 tab_batch = gr.Interface(
  fn=predict,
- inputs=gr.File(label="Upload multiple images", type="filepath", file_count="multiple"),
+ inputs=[
+ gr.File(label="Upload multiple images", type="filepath", file_count="multiple"),
+ gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`.", label="Resolution"),
+ gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.")
+ ],
  outputs=[gr.Gallery(label="BiRefNet's predictions"), gr.File(label="Download masked images.")],
  api_name="batch",
  description=descriptions+'\nTab-batch is partially modified from https://huggingface.co/spaces/NegiTurkey/Multi_Birefnetfor_Background_Removal, thanks to this great work!',

app_local.py

@@ -61,7 +61,7 @@ def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
 class ImagePreprocessor():
  def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
  self.transform_image = transforms.Compose([
- transforms.Resize(resolution), # 1. keep consistent with the cv2.resize used in training 2. redundant with that in path_to_image()
+ transforms.Resize(resolution),
  transforms.ToTensor(),
  transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
  ])
@@ -74,6 +74,7 @@ class ImagePreprocessor():
 usage_to_weights_file = {
  'General': 'BiRefNet',
  'General-Lite': 'BiRefNet_lite',
+ 'General-Lite-2K': 'BiRefNet_lite-2K',
  'Portrait': 'BiRefNet-portrait',
  'DIS': 'BiRefNet-DIS5K',
  'HRSOD': 'BiRefNet-HRSOD',
@@ -102,8 +103,8 @@ def predict(images, resolution, weights_file):
  try:
  resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
  except:
- resolution = [1024, 1024]
- print('Invalid resolution input. Automatically changed to 1024x1024.')
+ resolution = (1024, 1024) if weights_file not in ['General-Lite-2K'] else (2560, 1440)
+ print('Invalid resolution input. Automatically changed to 1024x1024 or 2K.')
 
  if isinstance(images, list):
  # For tab_batch
@@ -182,7 +183,7 @@ tab_image = gr.Interface(
  fn=predict,
  inputs=[
  gr.Image(label='Upload an image'),
- gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`. Higher resolutions can be much slower for inference.", label="Resolution"),
+ gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`.", label="Resolution"),
  gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.")
  ],
  outputs=gr.Image(label="BiRefNet's prediction", type="pil", format='png'),
@@ -195,7 +196,7 @@ tab_text = gr.Interface(
  fn=predict,
  inputs=[
  gr.Textbox(label="Paste an image URL"),
- gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`. Higher resolutions can be much slower for inference.", label="Resolution"),
+ gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`.", label="Resolution"),
  gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.")
  ],
  outputs=gr.Image(label="BiRefNet's prediction", type="pil", format='png'),
@@ -206,7 +207,11 @@ tab_text = gr.Interface(
 
 tab_batch = gr.Interface(
  fn=predict,
- inputs=gr.File(label="Upload multiple images", type="filepath", file_count="multiple"),
+ inputs=[
+ gr.File(label="Upload multiple images", type="filepath", file_count="multiple"),
+ gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`.", label="Resolution"),
+ gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.")
+ ],
  outputs=[gr.Gallery(label="BiRefNet's predictions"), gr.File(label="Download masked images.")],
  api_name="batch",
  description=descriptions+'\nTab-batch is partially modified from https://huggingface.co/spaces/NegiTurkey/Multi_Birefnetfor_Background_Removal, thanks to this great work!',