update app.py

app.py

@@ -5,7 +5,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import gradio as gr
 import codecs
-import numpy as np 
+import numpy as np
 import cv2
 
 from PIL import Image
@@ -15,41 +15,45 @@ from torchvision import transforms as T
 
 NORMALIZE = T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
 
-if __name__ == "__main__":
-    
-    def inference(img_path):
-        # Load the image
-        with open(img_path, "rb") as f:
-            img = Image.open(f)
-            img = img.convert("RGB")
-            img_np = np.array(img)
 
-            # Preprocess
-            t = T.Compose([T.ToTensor(), NORMALIZE])
-            img_t = t(img)[None, :, :, :]
-            inputs = img_t.to(device)
+def inference(img_path):
+    # Load the image
+    with open(img_path, "rb") as f:
+        img = Image.open(f)
+        img = img.convert("RGB")
+        img_np = np.array(img)
+
+        # Preprocess
+        t = T.Compose([T.ToTensor(), NORMALIZE])
+        img_t = t(img)[None, :, :, :]
+        inputs = img_t.to(device)
+
+    # Forward step
+    print(f"Start Peekaboo prediction.")
+    with torch.no_grad():
+        preds = model(inputs, for_eval=True)
+    print(f"Done Peekaboo prediction.")
 
-        # Forward step
-        print(f"Start Peekaboo prediction.")
-        with torch.no_grad():
-            preds = model(inputs, for_eval=True)
-        print(f"Done Peekaboo prediction.")
+    sigmoid = nn.Sigmoid()
+    h, w = img_t.shape[-2:]
+    preds_up = F.interpolate(
+        preds,
+        scale_factor=model.vit_patch_size,
+        mode="bilinear",
+        align_corners=False,
+    )[..., :h, :w]
+    preds_up = (sigmoid(preds_up.detach()) > 0.5).squeeze(0).float()
+    preds_up = preds_up.cpu().squeeze().numpy()
 
-        sigmoid = nn.Sigmoid()
-        h, w = img_t.shape[-2:]
-        preds_up = F.interpolate(
-            preds, scale_factor=model.vit_patch_size, mode="bilinear", align_corners=False
-        )[..., :h, :w]
-        preds_up = (sigmoid(preds_up.detach()) > 0.5).squeeze(0).float()
-        preds_up = preds_up.cpu().squeeze().numpy()
+    # Overlay predicted mask with input image
+    preds_up_np = (preds_up / np.max(preds_up) * 255).astype(np.uint8)
+    preds_up_np_3d = np.stack([preds_up_np, preds_up_np, preds_up_np], axis=-1)
+    combined_image = cv2.addWeighted(img_np, 0.5, preds_up_np_3d, 0.5, 0)
+    print(f"Output shape is {combined_image.shape}")
+    return combined_image
 
-        # Overlay predicted mask with input image
-        preds_up_np = (preds_up / np.max(preds_up) * 255).astype(np.uint8)
-        preds_up_np_3d = np.stack([preds_up_np, preds_up_np, preds_up_np], axis=-1)
-        combined_image = cv2.addWeighted(img_np, 0.5, preds_up_np_3d, 0.5, 0)
-        print(f"Output shape is {combined_image.shape}")
-        return combined_image
-    
+
+if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description="Evaluation of Peekaboo",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
@@ -108,7 +112,7 @@ if __name__ == "__main__":
             "./data/examples/b.jpeg",
             "./data/examples/c.jpeg",
             "./data/examples/d.jpeg",
-            "./data/examples/e.jpeg"
+            "./data/examples/e.jpeg",
         ],
         title=title,
         description=description,