Update app.py

app.py

@@ -15,6 +15,19 @@ model.load_state_dict(torch.load('pretrained_models/TranSalNet_Res.pth', map_loc
 model.to(device)
 model.eval()
 
+def count_and_label_red_patches(heatmap, threshold=200):
+    red_mask = heatmap[:, :, 2] > threshold
+    _, labels, stats, _ = cv2.connectedComponentsWithStats(red_mask.astype(np.uint8), connectivity=8)
+
+    num_red_patches = labels.max()
+    
+    for i in range(1, num_red_patches + 1):
+        patch_mask = (labels == i)
+        patch_centroid_x, patch_centroid_y = int(stats[i, cv2.CC_STAT_LEFT] + stats[i, cv2.CC_STAT_WIDTH] / 2), int(stats[i, cv2.CC_STAT_TOP] + stats[i, cv2.CC_STAT_HEIGHT] / 2)
+        cv2.putText(heatmap, str(i), (patch_centroid_x, patch_centroid_y), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2, cv2.LINE_AA)
+    
+    return heatmap, num_red_patches
+
 st.title('Saliency Detection App')
 st.write('Upload an image for saliency detection:')
 uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
@@ -25,57 +38,33 @@ if uploaded_image:
 
     if st.button('Detect Saliency'):
         img = image.resize((384, 288))
-        img = np.array(img) / 100.
+        img = np.array(img) / 255.
         img = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)
         img = torch.from_numpy(img)
         img = img.type(torch.FloatTensor).to(device)
 
-        pred_saliency = model(img)
+        pred_saliency = model(img).squeeze().detach().numpy()
 
-        toPIL = transforms.ToPILImage()
-        pic = toPIL(pred_saliency.squeeze())
+        heatmap = (pred_saliency * 255).astype(np.uint8)
+        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) 
 
-        colorized_img = cv2.applyColorMap(np.uint8(pic), cv2.COLORMAP_OCEAN)
+        heatmap = cv2.resize(heatmap, (image.width, image.height))
 
-        original_img = np.array(image)
-        colorized_img = cv2.resize(colorized_img, (original_img.shape[1], original_img.shape[0]))
+        heatmap, num_red_patches = count_and_label_red_patches(heatmap)
 
-        alpha = 0.7  
-        blended_img = cv2.addWeighted(original_img, 1 - alpha, colorized_img, alpha, 0)
+        enhanced_image = np.array(image)
+        b, g, r = cv2.split(enhanced_image)
+        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+        b_enhanced = clahe.apply(b)
+        enhanced_image = cv2.merge((b_enhanced, g, r))
 
-        # Find all contours
-        contours, _ = cv2.findContours(np.uint8(pred_saliency.squeeze().detach().numpy() * 255), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        saliency_8bit = np.uint8(pred_saliency.squeeze().detach().numpy() * 255)
-    
-        # Apply dilation
-        kernel = np.ones((5,5),np.uint8)
-        dilated = cv2.dilate(saliency_8bit, kernel, iterations = 1)
-    
-        # Find contours on dilated image
-        contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    
-        font = cv2.FONT_HERSHEY_SIMPLEX
-        label = 1
-        for contour in contours:
-            # Get bounding box for contour
-            x, y, w, h = cv2.boundingRect(contour)
-    
-            # Calculate center of bounding box
-            center_x = x + w // 2
-            center_y = y + h // 2
-    
-            # Find point on contour closest to center of bounding box
-            distances = np.sqrt((contour[:,0,0] - center_x)**2 + (contour[:,0,1] - center_y)**2)
-            min_index = np.argmin(distances)
-            closest_point = tuple(contour[min_index][0])
-    
-            # Place label at closest point on contour
-            cv2.putText(blended_img, str(label), closest_point, font, 1, (0, 0, 255), 3, cv2.LINE_AA)
-    
-            label += 1
+        alpha = 0.7 
+        blended_img = cv2.addWeighted(enhanced_image, 1 - alpha, heatmap, alpha, 0)
 
+        st.image(heatmap, caption='Enhanced Saliency Heatmap', use_column_width=True, channels='BGR')
+        st.image(enhanced_image, caption='Enhanced Blue Image', use_column_width=True, channels='BGR')
 
-        st.image(blended_img, caption='Blended Image with Labels', use_column_width=True)
+        st.image(blended_img, caption=f'Blended Image with {num_red_patches} Red Patches', use_column_width=True, channels='BGR')
 
         cv2.imwrite('example/result15.png', blended_img, [int(cv2.IMWRITE_JPEG_QUALITY), 200])
         st.success('Saliency detection complete. Result saved as "example/result15.png".')