Upload folder using huggingface_hub

app.py

@@ -40,7 +40,7 @@ load_model_without_module(sod_model, local_model_path)
 
 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}
-lrs = np.logspace(np.log10(0.0025), np.log10(0.25), 10).tolist()
+lrs = np.logspace(np.log10(0.001), np.log10(0.1), 10).tolist()
 img_size = 512
 
 cached_style_features = {}
@@ -54,7 +54,8 @@ for style_name, style_img_path in style_options.items():
 def run(content_image, style_name, style_strength=10):
     yield [None] * 3
     content_img, original_size = preprocess_img(content_image, img_size)
-    content_img = content_img.to(device)
+    content_img_normalized, _ = preprocess_img(content_image, img_size, normalize=True)
+    content_img, content_img_normalized = content_img.to(device), content_img_normalized.to(device)
     
     print('-'*15)
     print('DATETIME:', datetime.now(timezone.utc) - timedelta(hours=4)) # est
@@ -79,6 +80,7 @@ def run(content_image, style_name, style_strength=10):
             model=model,
             sod_model=sod_model,
             content_image=content_img,
+            content_image_norm=content_img_normalized,
             style_features=style_features,
             lr=lrs[style_strength-1],
             apply_to_background=apply_to_background

inference.py

@@ -38,6 +38,7 @@ def inference(
     model,
     sod_model,
     content_image,
+    content_image_norm,
     style_features,
     apply_to_background,
     lr,
@@ -56,7 +57,7 @@ def inference(
         resized_bg_masks = []
         salient_object_ratio = None
         if apply_to_background:
-            segmentation_output = sod_model(content_image)[0]
+            segmentation_output = sod_model(content_image_norm)[0]
             segmentation_output = torch.sigmoid(segmentation_output)
             segmentation_mask = (segmentation_output > 0.7).float()
             background_mask = (segmentation_mask == 0).float()

utils.py

@@ -3,28 +3,35 @@ from PIL import Image
 import torch
 import torchvision.transforms as transforms
 
-def preprocess_img_from_path(path_to_image, img_size):
+def preprocess_img_from_path(path_to_image, img_size, normalize=False):
     img = Image.open(path_to_image)
-    return preprocess_img(img, img_size)
+    return preprocess_img(img, img_size, normalize)
 
-def preprocess_img(img: Image, img_size):
+def preprocess_img(img: Image, img_size, normalize=False):
     original_size = img.size
     
-    transform = transforms.Compose([
-        transforms.Resize((img_size, img_size)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-    ])
+    if normalize:
+        transform = transforms.Compose([
+            transforms.Resize((img_size, img_size)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+    else:
+        transform = transforms.Compose([
+            transforms.Resize((img_size, img_size)),
+            transforms.ToTensor()
+        ])
     img = transform(img).unsqueeze(0)
     return img, original_size
 
-def postprocess_img(img, original_size):
+def postprocess_img(img, original_size, normalize=False):
     img = img.detach().cpu().squeeze(0)
     
     # Denormalize the image
-    mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
-    std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
-    img = img * std + mean
+    if normalize:
+        mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
+        std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
+        img = img * std + mean
     img = torch.clamp(img, 0, 1)
     
     img = transforms.ToPILImage()(img)