Update app.py

app.py

@@ -21,7 +21,6 @@ os.environ["HF_MODULES_CACHE"] = os.path.join("/tmp/hf_cache", "modules")
 import transformers
 transformers.utils.move_cache()
 
-
 torch.set_float32_matmul_precision('high')
 torch.jit.script = lambda f: f
 
@@ -37,15 +36,11 @@ def refine_foreground(image, mask, r=90):
     image_masked = Image.fromarray((estimated_foreground * 255.0).astype(np.uint8))
     return image_masked
 
-
 def FB_blur_fusion_foreground_estimator_2(image, alpha, r=90):
-    # Thanks to the source: https://github.com/Photoroom/fast-foreground-estimation
     alpha = alpha[:, :, None]
-    F, blur_B = FB_blur_fusion_foreground_estimator(
-        image, image, image, alpha, r)
+    F, blur_B = FB_blur_fusion_foreground_estimator(image, image, image, alpha, r)
     return FB_blur_fusion_foreground_estimator(image, F, blur_B, alpha, r=6)[0]
 
-
 def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
     if isinstance(image, Image.Image):
         image = np.array(image) / 255.0
@@ -56,15 +51,12 @@ def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
 
     blurred_B1A = cv2.blur(B * (1 - alpha), (r, r))
     blurred_B = blurred_B1A / ((1 - blurred_alpha) + 1e-5)
-    F = blurred_F + alpha * \
-        (image - alpha * blurred_F - (1 - alpha) * blurred_B)
+    F = blurred_F + alpha * (image - alpha * blurred_F - (1 - alpha) * blurred_B)
     F = np.clip(F, 0, 1)
     return F, blurred_B
 
-
 class ImagePreprocessor():
     def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
-        # Input resolution is on WxH.
         self.transform_image = transforms.Compose([
             transforms.Resize(resolution[::-1]),
             transforms.ToTensor(),
@@ -72,9 +64,7 @@ class ImagePreprocessor():
         ])
 
     def proc(self, image: Image.Image) -> torch.Tensor:
-        image = self.transform_image(image)
-        return image
-
+        return self.transform_image(image)
 
 usage_to_weights_file = {
     'General': 'BiRefNet',
@@ -94,17 +84,18 @@ usage_to_weights_file = {
     'General-dynamic': 'BiRefNet_dynamic',
 }
 
-birefnet = transformers.AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True)
+birefnet = transformers.AutoModelForImageSegmentation.from_pretrained(
+    '/'.join(('zhengpeng7', usage_to_weights_file['General'])),
+    trust_remote_code=True
+)
 birefnet.to(device)
 birefnet.eval(); birefnet.half()
 
-
 @spaces.GPU
 def predict(images, resolution, weights_file):
-    assert (images is not None), 'AssertionError: images cannot be None.'
+    assert images is not None, 'AssertionError: images cannot be None.'
 
     global birefnet
-    # Load BiRefNet with chosen weights
     _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: {}.'.format(_weights_file))
     birefnet = transformers.AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True)
@@ -129,7 +120,6 @@ def predict(images, resolution, weights_file):
         print('Invalid resolution input. Automatically changed to 1024x1024 / 2048x2048 / 2560x1440.')
 
     if isinstance(images, list):
-        # For tab_batch
         save_paths = []
         save_dir = 'preds-BiRefNet'
         if not os.path.exists(save_dir):
@@ -151,21 +141,17 @@ def predict(images, resolution, weights_file):
             image_ori = Image.fromarray(image_src)
 
         image = image_ori.convert('RGB')
-        # Preprocess the image
         if resolution is None:
             resolution_div_by_32 = [int(int(reso)//32*32) for reso in image.size]
             if resolution_div_by_32 != resolution:
                 resolution = resolution_div_by_32
         image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
-        image_proc = image_preprocessor.proc(image)
-        image_proc = image_proc.unsqueeze(0)
+        image_proc = image_preprocessor.proc(image).unsqueeze(0)
 
-        # Prediction
         with torch.no_grad():
             preds = birefnet(image_proc.to(device).half())[-1].sigmoid().cpu()
         pred = preds[0].squeeze()
 
-        # Show Results
         pred_pil = transforms.ToPILImage()(pred)
         image_masked = refine_foreground(image, pred_pil)
         image_masked.putalpha(pred_pil.resize(image.size))
@@ -184,32 +170,13 @@ def predict(images, resolution, weights_file):
                 zipf.write(file, os.path.basename(file))
         return save_paths, zip_file_path
     else:
-        return (image_masked, image_ori)
-
-
-examples = [[_] for _ in glob('examples/*')][:]
-# Add the option of resolution in a text box.
-for idx_example, example in enumerate(examples):
-    if 'My_' in example[0]:
-        example_resolution = '2048x2048'
-    else:
-        example_resolution = '1024x1024'
-    examples[idx_example].append(example_resolution)
-examples.append(examples[-1].copy())
-examples[-1][1] = '512x512'
-
-examples_url = [
-    ['https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg'],
-]
-for idx_example_url, example_url in enumerate(examples_url):
-    examples_url[idx_example_url].append('1024x1024')
+        return image_masked, image_ori
 
 descriptions = (
     "Upload a picture, and we'll remove the background!\n"
     "The resolution used is `1024x1024`\n"
 )
 
-                 
 tab_image = gr.Interface(
     fn=predict,
     inputs=[
@@ -218,7 +185,6 @@ tab_image = gr.Interface(
         gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.")
     ],
     outputs=gr.ImageSlider(label="BiRefNet's prediction", type="pil", format='png'),
-    examples=examples,
     api_name="image",
     description=descriptions,
 )
@@ -231,9 +197,8 @@ tab_text = gr.Interface(
         gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.")
     ],
     outputs=gr.ImageSlider(label="BiRefNet's prediction", type="pil", format='png'),
-    examples=examples_url,
     api_name="URL",
-    description=descriptions+'\nTab-URL is partially modified from https://huggingface.co/spaces/not-lain/background-removal, thanks to this great work!',
+    description=descriptions + '\nTab-URL is partially modified from https://huggingface.co/spaces/not-lain/background-removal, thanks to this great work!',
 )
 
 tab_batch = gr.Interface(
@@ -245,7 +210,7 @@ tab_batch = gr.Interface(
     ],
     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!',
+    description=descriptions + '\nTab-batch is partially modified from https://huggingface.co/spaces/NegiTurkey/Multi_Birefnetfor_Background_Removal, thanks to this great work!',
 )
 
 demo = gr.TabbedInterface(