Upload folder using huggingface_hub

inference.py

@@ -1,16 +1,9 @@
 import torch
 import torch.optim as optim
 import torch.nn.functional as F
-import matplotlib.pyplot as plt
+from torch.utils.tensorboard import SummaryWriter
 from torchvision.transforms.functional import gaussian_blur
 
-def save_mask(mask, title='mask'):
-    plt.imshow(mask.cpu().numpy()[0], cmap='gray')
-    plt.title(title)
-    plt.axis('off')
-    plt.savefig(f'{title}.png', bbox_inches='tight')
-    plt.close()
-
 def _gram_matrix(feature):
     batch_size, n_feature_maps, height, width = feature.size()
     new_feature = feature.view(batch_size * n_feature_maps, height * width)
@@ -35,7 +28,8 @@ def _compute_loss(generated_features, content_features, style_features, resized_
             A = _gram_matrix(sf)
         style_loss += w_l * F.mse_loss(G, A)
         
-    return alpha * content_loss + beta * style_loss
+    total_loss = alpha * content_loss + beta * style_loss
+    return content_loss, style_loss, total_loss
 
 def inference(
     *,
@@ -50,6 +44,7 @@ def inference(
     alpha=1,
     beta=1,
 ):
+    writer = SummaryWriter()
     generated_image = content_image.clone().requires_grad_(True)
     optimizer = optim_caller([generated_image], lr=lr)
     min_losses = [float('inf')] * iterations
@@ -64,12 +59,10 @@ def inference(
             segmentation_mask = segmentation_output.argmax(dim=1)
             background_mask = (segmentation_mask == 0).float()
             foreground_mask = 1 - background_mask
-            save_mask(background_mask, title='background-mask')
             
             background_pixel_count = background_mask.sum().item()
             total_pixel_count = segmentation_mask.numel()
             background_ratio = background_pixel_count / total_pixel_count
-            print(f'Background Detected: {background_ratio * 100:.2f}%')
 
             for cf in content_features:
                 _, _, h_i, w_i = cf.shape
@@ -79,11 +72,19 @@ def inference(
     def closure(iter):
         optimizer.zero_grad()
         generated_features = model(generated_image)
-        total_loss = _compute_loss(
+        content_loss, style_loss, total_loss = _compute_loss(
             generated_features, content_features, style_features, resized_bg_masks, alpha, beta
         )
         total_loss.backward()
+        
+        # log loss
+        writer.add_scalars(f'style-{"background" if apply_to_background else "image"}', {
+            'Loss/content': content_loss.item(),
+            'Loss/style': style_loss.item(),
+            'Loss/total': total_loss.item()
+        }, iter)
         min_losses[iter] = min(min_losses[iter], total_loss.item())
+        
         return total_loss
     
     for iter in range(iterations):
@@ -94,6 +95,6 @@ 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
 
-        if iter % 10 == 0: print(f'[{"Background" if apply_to_background else "Image"}] Loss ({iter}):', min_losses[iter])
-
+    writer.flush()
+    writer.close()
     return generated_image, background_ratio

u2net.py

@@ -1 +1,6 @@
-# initial code for u2net arch
+import torch.nn as nn
+
+class U2Net(nn.Module):
+    def __init__(self):
+        super(U2Net, self).__init__()
+        pass