Update run.py

run.py

@@ -4,38 +4,46 @@ import numpy as np
 import matplotlib.pyplot as plt
 import app
 import os
-import PIL
+import json
+from PIL import Image
 
 def pred(src):
-    # os.environ['HUGGINGFACE_HUB_HOME'] = './.cache'
-    # Load the model configuration
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model.to(device)
-    
-    cache_dir = "/code/cache"
-    model_config = SamConfig.from_pretrained("facebook/sam-vit-base", 
-                                             cache_dir=cache_dir)
-    processor = SamProcessor.from_pretrained("facebook/sam-vit-base", 
-                                             cache_dir=cache_dir)
 
-    # Create an instance of the model architecture with the loaded configuration
+    # -- load model configuration
+    MODEL_FILE = "sam_model.pth"
+    model_config = SamConfig.from_pretrained("facebook/sam-vit-base")
+    processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
+
     model = SamModel(config=model_config)
-    #Update the model by loading the weights from saved file.
-    model.load_state_dict(torch.load("sam_model.pth", 
-                                     map_location=torch.device('cpu')))
+    model.load_state_dict(torch.load(MODEL_FILE))
+
+    with open("sam-config.json", "r") as f: # modified config json file
+        modified_config_dict = json.load(f)
+
+    processor = SamProcessor.from_pretrained("facebook/sam-vit-base", 
+                                            **modified_config_dict)
+    
+    # -- process image
+    image = Image.open(src)
+    rgbim = image.convert("RGB")
+    new_image = np.array(rgbim)
+    print("Shape:",new_image.shape)
 
-    new_image = np.array(Image.open(src))
     inputs = processor(new_image, return_tensors="pt")
-    inputs = {k: v.to(device) for k, v in inputs.items()}
-    x = 1
-    # model.eval()
-    # # forward pass
-    # with torch.no_grad():
-    #     outputs = model(**inputs, multimask_output=False)
-
-    # # apply sigmoid
-    # single_patch_prob = torch.sigmoid(outputs.pred_masks.squeeze(1))
-    # # convert soft mask to hard mask
-    # single_patch_prob = single_patch_prob.cpu().numpy().squeeze()
-    # single_patch_prediction = (single_patch_prob > 0.5).astype(np.uint8)
+    model.eval()
+
+    # forward pass
+    with torch.no_grad():
+        outputs = model(pixel_values=inputs["pixel_values"], 
+                        multimask_output=False)
+    
+    # apply sigmoid
+    pred_prob = torch.sigmoid(outputs.pred_masks.squeeze(1))
+
+    # convert soft mask to hard mask
+    PROBABILITY_THRES = 0.30
+    pred_prob = pred_prob.cpu().numpy().squeeze()
+    pred_prediction = (pred_prob > PROBABILITY_THRES).astype(np.uint8)
+    
+    x=1
     return x