Update app.py

app.py

@@ -68,35 +68,89 @@ def build_model(hypar,device):
     net.eval()  
     return net
 
-    
-def predict(net,  inputs_val, shapes_val, hypar, device):
-    '''
-    Given an Image, predict the mask
-    '''
-    net.eval()
-
-    if(hypar["model_digit"]=="full"):
-        inputs_val = inputs_val.type(torch.FloatTensor)
+def resize_image(image, size=1024):
+
+    height, width = image.shape[:2]
+
+    # Check if either dimension is greater than 1120
+    if height > size or width > size:
+        # Calculate the scale factor
+        if height > width:
+            scale_factor = size / height
+        else:
+            scale_factor = size / width
+
+        # Resize the image
+        new_dimensions = (int(width * scale_factor), int(height * scale_factor))
+        resized_image = cv2.resize(image, new_dimensions, interpolation=cv2.INTER_AREA)
+        
+        # Save the resized image
+        
+        print(f"Image resized to {new_dimensions}")
+        return resized_image
     else:
-        inputs_val = inputs_val.type(torch.HalfTensor)
-
-  
-    inputs_val_v = Variable(inputs_val, requires_grad=False).to(device) # wrap inputs in Variable
-   
-    ds_val = net(inputs_val_v)[0] # list of 6 results
-
-    pred_val = ds_val[0][0,:,:,:] # B x 1 x H x W    # we want the first one which is the most accurate prediction
-
-    ## recover the prediction spatial size to the orignal image size
-    pred_val = torch.squeeze(F.upsample(torch.unsqueeze(pred_val,0),(shapes_val[0][0],shapes_val[0][1]),mode='bilinear'))
-
-    ma = torch.max(pred_val)
-    mi = torch.min(pred_val)
-    pred_val = (pred_val-mi)/(ma-mi) # max = 1
+        print("Image is already within the desired size.")
+        return image
 
-    if device == 'cuda': torch.cuda.empty_cache()
-    return (pred_val.detach().cpu().numpy()*255).astype(np.uint8) # it is the mask we need
+def predict(net,  im):
     
+    im = resize_image(im)
+    temp = np.ones((1024,1024,3))
+    h, w = im.shape[0],im.shape[1]
+    temp[:h,:w] = im
+    im = temp
+    #show_pic(im)
+    input_size = [1024,1024]
+    if len(im.shape) < 3:
+        im = np.stack([im] * 3, axis=-1)  # Convert grayscale to RGB
+    im_shp = im.shape[0:2]
+    im_tensor = torch.tensor(im, dtype=torch.float32).permute(2, 0, 1)
+    im_tensor = F.upsample(torch.unsqueeze(im_tensor, 0), input_size, mode="bilinear").type(torch.uint8)
+    image = torch.divide(im_tensor, 255.0)
+    image = normalize(image, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
+
+    result = net(image)
+    result = torch.squeeze(F.upsample(result[0][0], im_shp, mode='bilinear'), 0)
+    ma = torch.max(result)
+    mi = torch.min(result)
+    result = (result - mi) / (ma - mi)
+    result = result.unsqueeze(0) if result.dim() == 2 else result  # Ensure result has 3 channels
+    result = result.repeat(3, 1, 1) if result.shape[0] == 1 else result
+    result = 1 - result  # Invert the mask here
+
+
+
+    #im_name = im_path.split('\\')[-1].split('.')[0]
+
+    # Resize the image to match result dimensions
+    image_resized = F.upsample(image, size=result.shape[1:], mode='bilinear')
+
+    # Ensure both tensors are 3D
+    image_resized = image_resized.squeeze(0) if image_resized.dim() == 4 else image_resized
+    result = result.squeeze(0) if result.dim() == 4 else result
+
+    # Apply threshold to result to ensure only pure black or white pixels
+    threshold = 0.50  # Adjust as needed
+    result[result < threshold] = 0
+    result[result >= threshold] = 1
+
+    distance = np.sqrt(np.sum((im - [255, 255, 255]) ** 2, axis=-1))
+
+    # Create a mask where the distance is less than the threshold
+    mask = distance < 200
+
+    # Convert mask to uint8
+    mask = mask.astype(np.uint8) * 255
+
+    mask = np.stack([mask] * 3, axis=-1)
+
+    result = (result.permute(1, 2, 0) * 255).cpu().numpy().astype(np.uint8)
+    # result=result.cpu().numpy().astype(np.uint8)
+    # io.imsave(result_path + im_name + "_foreground.png", foreground)
+    wite = np.ones_like(im) * 255
+    cropped = np.where(result == 0, wite, mask)
+    #cv2.imwrite(result_path + f, cropped)
+    return cropped[:h,:w]
 # Set Parameters
 hypar = {} # paramters for inferencing
 
@@ -124,16 +178,11 @@ net = build_model(hypar, device)
 def inference(image):
   image_path = image
   
-  image_tensor, orig_size = load_image(image_path, hypar) 
-  mask = predict(net, image_tensor, orig_size, hypar, device)
-  
-  pil_mask = Image.fromarray(mask).convert('L')
-  im_rgb = Image.open(image).convert("RGB")
+  image_tensor, orig_size = cv2.imread(image_path) 
+  mask = predict(net, image_tensor)
   
-  im_rgba = im_rgb.copy()
-  im_rgba.putalpha(pil_mask)
 
-  return [im_rgba, pil_mask]
+  return [mask,mask]
 
 
 title = "Highly Accurate Dichotomous Image Segmentation"