Update app.py

app.py

@@ -16,32 +16,45 @@ 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()
-
+    contours, _ = cv2.findContours(red_mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    
+    # Sort the contours based on their areas in descending order
+    contours = sorted(contours, key=cv2.contourArea, reverse=True)
+    
     original_image = np.array(image)
-
-    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)
-        radius = 20  # Adjust the following variable to manage the circle image
-        circle_color = (0, 0, 0)  # The circle is black adjust the following to change the color
-        cv2.circle(original_image, (patch_centroid_x, patch_centroid_y), radius, circle_color, -1)  # Draw the circle
-
-        # Lines code
-        for j in range(i + 1, num_red_patches + 1):
-            patch_mask_j = (labels == j)
-            patch_centroid_x_j, patch_centroid_y_j = int(stats[j, cv2.CC_STAT_LEFT] + stats[j, cv2.CC_STAT_WIDTH] / 2), int(stats[j, cv2.CC_STAT_TOP] + stats[j, cv2.CC_STAT_HEIGHT] / 2)
-            line_color = (0, 0, 0)  # Ajdust the following to manage the line color
-            cv2.line(original_image, (patch_centroid_x, patch_centroid_y), (patch_centroid_x_j, patch_centroid_y_j), line_color, 2) # Line
+    
+    # Find the centroid of the red spot with the highest area
+    M_largest = cv2.moments(contours[0])
+    if M_largest["m00"] != 0:
+        cX_largest = int(M_largest["m10"] / M_largest["m00"])
+        cY_largest = int(M_largest["m01"] / M_largest["m00"])
+    else:
+        cX_largest, cY_largest = 0, 0
+    
+    for i, contour in enumerate(contours, start=1):
+        # Compute the centroid of the current contour
+        M = cv2.moments(contour)
+        if M["m00"] != 0:
+            cX = int(M["m10"] / M["m00"])
+            cY = int(M["m01"] / M["m00"])
+        else:
+            cX, cY = 0, 0
+        
+        radius = 20  # Adjust the circle radius to fit the numbers
+        circle_color = (0, 0, 0)  # Blue color
+        cv2.circle(original_image, (cX, cY), radius, circle_color, -1)  # Draw blue circle
+
+        # Connect the current red spot to the red spot with the highest area
+        line_color = (0, 0, 0)  # Red color
+        cv2.line(original_image, (cX, cY), (cX_largest, cY_largest), line_color, 2)
 
         font = cv2.FONT_HERSHEY_SIMPLEX
         font_scale = 1
         font_color = (255, 255, 255)
         line_type = cv2.LINE_AA
-        cv2.putText(original_image, str(i), (patch_centroid_x - 10, patch_centroid_y + 10), font, font_scale, font_color, 2, line_type)
+        cv2.putText(original_image, str(i), (cX - 10, cY + 10), font, font_scale, font_color, 2, line_type)
 
-    return original_image, num_red_patches
+    return original_image, len(contours)
 
 st.title('Saliency Detection App')
 st.write('Upload an image for saliency detection:')
@@ -82,6 +95,6 @@ if uploaded_image:
 
         st.image(blended_img, caption='Blended Image', use_column_width=True, channels='BGR')
 
-        # Create a dir with the name example to save
+        # Create a dir with name example to save
         cv2.imwrite('example/result15.png', blended_img, [int(cv2.IMWRITE_JPEG_QUALITY), 200])
         st.success('Saliency detection complete. Result saved as "example/result15.png".')