Upload folder using huggingface_hub

app.py

@@ -33,11 +33,8 @@ def load_model_without_module(model, model_path):
 model = VGG_19().to(device).eval()
 for param in model.parameters():
     param.requires_grad = False
-# sod_model = U2Net().to(device).eval()
-# load_model_without_module(sod_model, 'u2net/saved_models/u2net-duts.pt')
-sod_model = models.segmentation.deeplabv3_resnet101(
-    weights='DEFAULT'
-).to(device).eval()
+sod_model = U2Net().to(device).eval()
+load_model_without_module(sod_model, 'u2net/saved_models/u2net-duts.pt')
 
 style_files = os.listdir('./style_images')
 style_options = {' '.join(style_file.split('.')[0].split('_')): f'./style_images/{style_file}' for style_file in style_files}
@@ -166,4 +163,4 @@ with gr.Blocks(css=css) as demo:
 
 demo.queue = False
 demo.config['queue'] = False
-demo.launch(show_api=False)
+demo.launch(show_api=False)

inference.py

@@ -56,19 +56,12 @@ def inference(
         resized_bg_masks = []
         salient_object_ratio = None
         if apply_to_background:
-            # original
-            segmentation_output = sod_model(content_image)['out'] # [1, 21, 512, 512]
-            segmentation_mask = segmentation_output.argmax(dim=1) # [1, 512, 512]
+            segmentation_output = sod_model(content_image)[0]
+            segmentation_output = torch.sigmoid(segmentation_output)
+            segmentation_mask = (segmentation_output > 0.7).float()
             background_mask = (segmentation_mask == 0).float()
             foreground_mask = 1 - background_mask
             
-            # new
-            # segmentation_output = sod_model(content_image)[0]
-            # segmentation_output = torch.sigmoid(segmentation_output)
-            # segmentation_mask = (segmentation_output > 0.7).float()
-            # background_mask = (segmentation_mask == 0).float()
-            # foreground_mask = 1 - background_mask
-            
             salient_object_pixel_count = foreground_mask.sum().item()
             total_pixel_count = segmentation_mask.numel()
             salient_object_ratio = salient_object_pixel_count / total_pixel_count
@@ -99,4 +92,4 @@ def inference(
                 foreground_mask_resized = F.interpolate(foreground_mask.unsqueeze(1), size=generated_image.shape[2:], mode='nearest')
                 generated_image.data = generated_image.data * (1 - foreground_mask_resized) + content_image.data * foreground_mask_resized
                 
-    return generated_image, salient_object_ratio
+    return generated_image, salient_object_ratio