Update app.py

app.py

@@ -108,8 +108,8 @@ def drawOnTop(img, landmarks, original_shape):
     image = cv2.line(image, (int(rl_top[0]), int(rl_top[1])), (int(ll_top[0]), int(ll_top[1])), (0, 1, 0), 1)
     
     # Add tilt angle text
-    tilt_text = f"Tilt: {tilt_angle:.1f}°"
-    cv2.putText(image, tilt_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 1, 0), 2)
+    tilt_text = f"Tilt: {tilt_angle:.1f} degrees"
+    cv2.putText(image, tilt_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 1, 0), 2)
     
     # Correct landmarks for tilt
     if abs(tilt_angle) > 2:  # Only correct if tilt is significant
@@ -120,51 +120,24 @@ def drawOnTop(img, landmarks, original_shape):
     else:
         RL_corrected, LL_corrected, H_corrected = RL, LL, H
 
-    # Get interpolated points for more precise measurements (MOVED UP)
-    # Reduced for HuggingFace memory constraints
-    H_dense = interpolate_contour_points(H_corrected, num_points=150)
-    RL_dense = interpolate_contour_points(RL_corrected, num_points=100)
-    LL_dense = interpolate_contour_points(LL_corrected, num_points=100)
-    
-    # Heart (red line) - calculate positions from corrected coordinates using precise method
-    heart_left, heart_right = find_precise_extremes(H_dense, direction='horizontal')
-    heart_y_corrected = np.mean([heart_left[1], heart_right[1]])
-    
-    # Thorax (blue line) - calculate positions from corrected coordinates using precise method
-    rl_left, rl_right = find_precise_extremes(RL_dense, direction='horizontal')
-    ll_left, ll_right = find_precise_extremes(LL_dense, direction='horizontal')
-    
-    # Get the overall leftmost and rightmost points
-    thorax_left = rl_left if rl_left[0] < ll_left[0] else ll_left
-    thorax_right = rl_right if rl_right[0] > ll_right[0] else ll_right
-    thorax_y_corrected = np.mean([thorax_left[1], thorax_right[1]])
-
-    # Add precision info (NOW AFTER CREATING THE VARIABLES)
-    precision_text = f"Enhanced: {len(H_dense)} heart pts, {len(RL_dense)+len(LL_dense)} lung pts"
-    cv2.putText(image, precision_text, (10, h-10), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (1, 1, 1), 1)
-
     # Draw the landmarks as dots
     for l in RL:
-        image = cv2.circle(image, (int(l[0]), int(l[1])), 3, (1, 0, 1), -1)
+        image = cv2.circle(image, (int(l[0]), int(l[1])), 5, (1, 0, 1), -1)
     for l in LL:
-        image = cv2.circle(image, (int(l[0]), int(l[1])), 3, (1, 0, 1), -1)
+        image = cv2.circle(image, (int(l[0]), int(l[1])), 5, (1, 0, 1), -1)
     for l in H:
-        image = cv2.circle(image, (int(l[0]), int(l[1])), 3, (1, 1, 0), -1)
-    
-    # Highlight the precise extreme points used for measurement
-    # Heart extreme points (larger red circles)
-    image = cv2.circle(image, (int(heart_left[0]), int(heart_left[1])), 8, (1, 0, 0), -1)
-    image = cv2.circle(image, (int(heart_right[0]), int(heart_right[1])), 8, (1, 0, 0), -1)
-    
-    # Thorax extreme points (larger blue circles)  
-    image = cv2.circle(image, (int(thorax_left[0]), int(thorax_left[1])), 8, (0, 0, 1), -1)
-    image = cv2.circle(image, (int(thorax_right[0]), int(thorax_right[1])), 8, (0, 0, 1), -1)
+        image = cv2.circle(image, (int(l[0]), int(l[1])), 5, (1, 1, 0), -1)
 
     # Draw measurement lines that follow the image tilt for visual accuracy
     # Use corrected coordinates for accurate measurement, but draw tilted lines for visual appeal
     
+    # Heart (red line) - calculate positions from corrected coordinates
+    heart_xmin_corrected = np.min(H_corrected[:, 0])
+    heart_xmax_corrected = np.max(H_corrected[:, 0])
+    heart_y_corrected = np.mean([H_corrected[np.argmin(H_corrected[:, 0]), 1], H_corrected[np.argmax(H_corrected[:, 0]), 1]])
+    
     # Rotate back to match the tilted image for display
-    heart_points_corrected = np.array([[heart_left[0], heart_y_corrected], [heart_right[0], heart_y_corrected]])
+    heart_points_corrected = np.array([[heart_xmin_corrected, heart_y_corrected], [heart_xmax_corrected, heart_y_corrected]])
     heart_points_display = rotate_points(heart_points_corrected, -tilt_angle, image_center)  # Rotate back for display
     
     heart_start = (int(heart_points_display[0, 0]), int(heart_points_display[0, 1]))
@@ -191,9 +164,24 @@ def drawOnTop(img, landmarks, original_shape):
                         (int(heart_end[0] + perp_x), int(heart_end[1] + perp_y)),
                         (int(heart_end[0] - perp_x), int(heart_end[1] - perp_y)), 
                         (1, 0, 0), 2)
+
+    # Thorax (blue line) - calculate positions from corrected coordinates
+    thorax_xmin_corrected = min(np.min(RL_corrected[:, 0]), np.min(LL_corrected[:, 0]))
+    thorax_xmax_corrected = max(np.max(RL_corrected[:, 0]), np.max(LL_corrected[:, 0]))
+    
+    # Find y at leftmost and rightmost points (corrected)
+    if np.min(RL_corrected[:, 0]) < np.min(LL_corrected[:, 0]):
+        thorax_ymin_corrected = RL_corrected[np.argmin(RL_corrected[:, 0]), 1]
+    else:
+        thorax_ymin_corrected = LL_corrected[np.argmin(LL_corrected[:, 0]), 1]
+    if np.max(RL_corrected[:, 0]) > np.max(LL_corrected[:, 0]):
+        thorax_ymax_corrected = RL_corrected[np.argmax(RL_corrected[:, 0]), 1]
+    else:
+        thorax_ymax_corrected = LL_corrected[np.argmax(LL_corrected[:, 0]), 1]
+    thorax_y_corrected = np.mean([thorax_ymin_corrected, thorax_ymax_corrected])
     
     # Rotate back to match the tilted image for display
-    thorax_points_corrected = np.array([[thorax_left[0], thorax_y_corrected], [thorax_right[0], thorax_y_corrected]])
+    thorax_points_corrected = np.array([[thorax_xmin_corrected, thorax_y_corrected], [thorax_xmax_corrected, thorax_y_corrected]])
     thorax_points_display = rotate_points(thorax_points_corrected, -tilt_angle, image_center)  # Rotate back for display
     
     thorax_start = (int(thorax_points_display[0, 0]), int(thorax_points_display[0, 1]))
@@ -305,162 +293,19 @@ def removePreprocess(output, info):
     return output
 
 
-def interpolate_contour_points(points, num_points=100):
-    """Interpolate additional points along a contour for more precise measurements"""
-    try:
-        if len(points) < 3:
-            return points
-        
-        # Reduce memory usage for HuggingFace
-        num_points = min(num_points, 150)  # Limit max points
-        
-        # Close the contour if not already closed
-        if not np.array_equal(points[0], points[-1]):
-            points = np.vstack([points, points[0]])
-        
-        # Calculate cumulative distances along the contour
-        distances = np.sqrt(np.sum(np.diff(points, axis=0)**2, axis=1))
-        cumulative_distances = np.concatenate([[0], np.cumsum(distances)])
-        
-        # Create new evenly spaced parameter values
-        total_length = cumulative_distances[-1]
-        if total_length == 0:
-            return points
-            
-        new_params = np.linspace(0, total_length, num_points)
-        
-        # Interpolate x and y coordinates
-        new_x = np.interp(new_params, cumulative_distances, points[:, 0])
-        new_y = np.interp(new_params, cumulative_distances, points[:, 1])
-        
-        result = np.column_stack([new_x, new_y]).astype(int)
-        return result
-    except Exception as e:
-        print(f"Error in interpolation: {e}")
-        return points
-
-def find_precise_extremes(points, direction='horizontal'):
-    """Find extreme points with sub-pixel precision using parabolic fitting"""
-    try:
-        if len(points) < 3:
-            return points[0], points[-1] if len(points) > 1 else points[0]
-            
-        if direction == 'horizontal':
-            # Find leftmost and rightmost points
-            coord_idx = 0  # x-coordinate
-        else:
-            # Find topmost and bottommost points
-            coord_idx = 1  # y-coordinate
-            
-        # Find approximate extremes
-        min_idx = np.argmin(points[:, coord_idx])
-        max_idx = np.argmax(points[:, coord_idx])
-        
-        def refine_extreme(idx, is_maximum=False):
-            """Refine extreme point using neighboring points"""
-            try:
-                if len(points) < 3:
-                    return points[idx]
-                    
-                # Get neighboring indices (with wrapping)
-                prev_idx = (idx - 1) % len(points)
-                next_idx = (idx + 1) % len(points)
-                
-                # Get the three points
-                p1 = points[prev_idx]
-                p2 = points[idx]
-                p3 = points[next_idx]
-                
-                # Simple parabolic fitting in the coordinate direction
-                coords = np.array([p1[coord_idx], p2[coord_idx], p3[coord_idx]])
-                if is_maximum:
-                    best_idx = np.argmax(coords)
-                else:
-                    best_idx = np.argmin(coords)
-                    
-                if best_idx == 1:  # Original point is still the best
-                    return p2
-                elif best_idx == 0:
-                    return p1
-                else:
-                    return p3
-            except Exception as e:
-                print(f"Error refining extreme: {e}")
-                return points[idx]
-        
-        # Refine the extreme points
-        min_point = refine_extreme(min_idx, is_maximum=False)
-        max_point = refine_extreme(max_idx, is_maximum=True)
-        
-        return min_point, max_point
-    except Exception as e:
-        print(f"Error finding extremes: {e}")
-        # Fallback to simple min/max
-        if direction == 'horizontal':
-            coord_idx = 0
-        else:
-            coord_idx = 1
-        min_idx = np.argmin(points[:, coord_idx])
-        max_idx = np.argmax(points[:, coord_idx])
-        return points[min_idx], points[max_idx]
-
 def calculate_ctr(landmarks, corrected_landmarks=None):
-    try:
-        if corrected_landmarks is not None:
-            RL, LL, H, tilt_angle = corrected_landmarks
-        else:
-            H = landmarks[94:]
-            RL = landmarks[0:44]
-            LL = landmarks[44:94]
-            tilt_angle = 0
-        
-        # Reduced points for HuggingFace
-        RL_dense = interpolate_contour_points(RL, num_points=100)
-        LL_dense = interpolate_contour_points(LL, num_points=100)
-        H_dense = interpolate_contour_points(H, num_points=150)
-        
-        # Find precise extreme points for heart
-        heart_left, heart_right = find_precise_extremes(H_dense, direction='horizontal')
-        cardiac_width = heart_right[0] - heart_left[0]
-        
-        # Find precise extreme points for thorax (combining both lungs)
-        rl_left, rl_right = find_precise_extremes(RL_dense, direction='horizontal')
-        ll_left, ll_right = find_precise_extremes(LL_dense, direction='horizontal')
-        
-        # Get the overall leftmost and rightmost points
-        thorax_left = rl_left if rl_left[0] < ll_left[0] else ll_left
-        thorax_right = rl_right if rl_right[0] > ll_right[0] else ll_right
-        thoracic_width = thorax_right[0] - thorax_left[0]
-        
-        # Calculate CTR with additional precision checks
-        if thoracic_width > 0 and cardiac_width > 0:
-            ctr = cardiac_width / thoracic_width
-            
-            # Sanity check: CTR should be between 0.2 and 1.0
-            if ctr < 0.2 or ctr > 1.0:
-                # Fallback to simple calculation
-                cardiac_width_simple = np.max(H[:, 0]) - np.min(H[:, 0])
-                thoracic_width_simple = max(np.max(RL[:, 0]), np.max(LL[:, 0])) - min(np.min(RL[:, 0]), np.min(LL[:, 0]))
-                ctr = cardiac_width_simple / thoracic_width_simple if thoracic_width_simple > 0 else 0
-        else:
-            ctr = 0
-        
-        return round(ctr, 3), abs(tilt_angle)
-    except Exception as e:
-        print(f"Error in CTR calculation: {e}")
-        # Simple fallback calculation
-        try:
-            H = landmarks[94:] if corrected_landmarks is None else corrected_landmarks[2]
-            RL = landmarks[0:44] if corrected_landmarks is None else corrected_landmarks[0]
-            LL = landmarks[44:94] if corrected_landmarks is None else corrected_landmarks[1]
-            tilt_angle = 0 if corrected_landmarks is None else corrected_landmarks[3]
-            
-            cardiac_width = np.max(H[:, 0]) - np.min(H[:, 0])
-            thoracic_width = max(np.max(RL[:, 0]), np.max(LL[:, 0])) - min(np.min(RL[:, 0]), np.min(LL[:, 0]))
-            ctr = cardiac_width / thoracic_width if thoracic_width > 0 else 0
-            return round(ctr, 3), abs(tilt_angle)
-        except:
-            return 0.5, 0
+    if corrected_landmarks is not None:
+        RL, LL, H, tilt_angle = corrected_landmarks
+    else:
+        H = landmarks[94:]
+        RL = landmarks[0:44]
+        LL = landmarks[44:94]
+        tilt_angle = 0
+    
+    cardiac_width = np.max(H[:, 0]) - np.min(H[:, 0])
+    thoracic_width = max(np.max(RL[:, 0]), np.max(LL[:, 0])) - min(np.min(RL[:, 0]), np.min(LL[:, 0]))
+    ctr = cardiac_width / thoracic_width if thoracic_width > 0 else 0
+    return round(ctr, 3), abs(tilt_angle)
 
 
 def detect_image_rotation(img):
@@ -527,79 +372,45 @@ def segment(input_img):
     global hybrid, device
 
     try:
-        # Validate input
-        if input_img is None:
-            return None, None, 0, "Error: No image provided"
-            
         if hybrid is None:
-            print("Loading model...")
             hybrid = loadModel(device)
-            print("Model loaded successfully")
 
-        # Load and validate image
-        original_img = cv2.imread(input_img, 0)
-        if original_img is None:
-            return None, None, 0, "Error: Could not load image"
-            
-        original_img = original_img / 255.0
+        original_img = cv2.imread(input_img, 0) / 255.0
         original_shape = original_img.shape[:2]
-        print(f"Image loaded: {original_shape}")
         
         # Step 1: For now, skip rotation detection to avoid errors
+        # TODO: Re-implement rotation detection after fixing coordinate transformation
         detected_rotation = 0  # Temporarily disabled
         was_rotated = False
         processing_img = original_img
 
         # Step 2: Preprocess the image
-        print("Preprocessing image...")
         img, (h, w, padding) = preprocess(processing_img)
 
         # Step 3: AI segmentation
-        print("Running AI segmentation...")
         data = torch.from_numpy(img).unsqueeze(0).unsqueeze(0).to(device).float()
 
         with torch.no_grad():
             output = hybrid(data)[0].cpu().numpy().reshape(-1, 2)
 
         # Step 4: Remove preprocessing
-        print("Post-processing...")
         output = removePreprocess(output, (h, w, padding))
         
         # Step 5: Convert output to int
         output = output.astype('int')
 
         # Step 6: Draw results on original image
-        print("Drawing results...")
         outseg, corrected_data = drawOnTop(original_img, output, original_shape)
         
-        if outseg is None:
-            return None, None, 0, "Error: Failed to generate output image"
-            
     except Exception as e:
         print(f"Error in segmentation: {e}")
-        import traceback
-        traceback.print_exc()
+        # Return a basic error response
         return None, None, 0, f"Error: {str(e)}"
 
-    try:
-        # Save output image
-        seg_to_save = (outseg.copy() * 255).astype('uint8')
-        
-        # Create tmp directory if it doesn't exist
-        import os
-        os.makedirs("tmp", exist_ok=True)
-        
-        success = cv2.imwrite("tmp/overlap_segmentation.png", cv2.cvtColor(seg_to_save, cv2.COLOR_RGB2BGR))
-        if not success:
-            print("Warning: Could not save output image")
-            
-        print("Calculating CTR...")
-        ctr_value, tilt_angle = calculate_ctr(output, corrected_data)
-        print(f"CTR calculated: {ctr_value}")
-        
-    except Exception as e:
-        print(f"Error in final processing: {e}")
-        return outseg, None, 0.5, f"Segmentation completed but error in final processing: {str(e)}"
+    seg_to_save = (outseg.copy() * 255).astype('uint8')
+    cv2.imwrite("tmp/overlap_segmentation.png", cv2.cvtColor(seg_to_save, cv2.COLOR_RGB2BGR))
+
+    ctr_value, tilt_angle = calculate_ctr(output, corrected_data)
     
     # Add rotation info to interpretation
     rotation_warning = ""
@@ -613,44 +424,34 @@ def segment(input_img):
     elif tilt_angle > 2:
         tilt_warning = f" (Minor tilt: {tilt_angle:.1f}°)"
     
-    # Enhanced interpretation with precision info (reduced for HuggingFace)
-    precision_note = " [Enhanced precision: 350+ interpolated points]"
-    
     if ctr_value < 0.5:
-        interpretation = f"Normal{rotation_warning}{tilt_warning}{precision_note}"
+        interpretation = f"Normal{rotation_warning}{tilt_warning}"
     elif 0.51 <= ctr_value <= 0.55:
-        interpretation = f"Mild Cardiomegaly (CTR 51-55%){rotation_warning}{tilt_warning}{precision_note}"
+        interpretation = f"Mild Cardiomegaly (CTR 51-55%){rotation_warning}{tilt_warning}"
     elif 0.56 <= ctr_value <= 0.60:
-        interpretation = f"Moderate Cardiomegaly (CTR 56-60%){rotation_warning}{tilt_warning}{precision_note}"
+        interpretation = f"Moderate Cardiomegaly (CTR 56-60%){rotation_warning}{tilt_warning}"
     elif ctr_value > 0.60:
-        interpretation = f"Severe Cardiomegaly (CTR > 60%){rotation_warning}{tilt_warning}{precision_note}"
+        interpretation = f"Severe Cardiomegaly (CTR > 60%){rotation_warning}{tilt_warning}"
     else:
-        interpretation = f"Cardiomegaly{rotation_warning}{tilt_warning}{precision_note}"
+        interpretation = f"Cardiomegaly{rotation_warning}{tilt_warning}"
 
     return outseg, "tmp/overlap_segmentation.png", ctr_value, interpretation
 
 
 if __name__ == "__main__":
-    # Clear any existing CUDA cache for HuggingFace
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        
-    with gr.Blocks(title="Chest X-ray Segmentation") as demo:
+    with gr.Blocks() as demo:
         gr.Markdown("""
-                    # Chest X-ray HybridGNet Segmentation
+                    # Chest X-ray HybridGNet Segmentation.
+
+                    Demo of the HybridGNet model introduced in "Improving anatomical plausibility in medical image segmentation via hybrid graph neural networks: applications to chest x-ray analysis."
 
-                    Enhanced CTR calculation with improved precision using interpolated landmarks.
+                    Instructions:
+                    1. Upload a chest X-ray image (PA or AP) in PNG or JPEG format.
+                    2. Click on "Segment Image".
 
-                    **Instructions:**
-                    1. Upload a chest X-ray image (PA or AP) in PNG or JPEG format
-                    2. Click "Segment Image" 
-                    3. Wait for processing (may take 30-60 seconds)
+                    Note: Pre-processing is not needed, it will be done automatically and removed after the segmentation.
 
-                    **Features:**
-                    - AI-powered lung and heart segmentation
-                    - Enhanced CTR calculation with 350+ interpolated points
-                    - Automatic tilt correction
-                    - Visual measurement lines with precision markers
+                    Please check citations below.                    
                     """)
 
         with gr.Tab("Segment Image"):
@@ -721,16 +522,6 @@ if __name__ == "__main__":
         clear_button.click(lambda: None, None, ctr_output, queue=False)
         clear_button.click(lambda: None, None, ctr_interpretation, queue=False)
 
-        image_button.click(
-            fn=segment, 
-            inputs=image_input, 
-            outputs=[image_output, results, ctr_output, ctr_interpretation],
-            show_progress=True
-        )
-
-    demo.launch(
-        share=False,
-        server_name="0.0.0.0",
-        server_port=7860,
-        enable_queue=True
-    )
+        image_button.click(segment, inputs=image_input, outputs=[image_output, results, ctr_output, ctr_interpretation], queue=False)
+
+    demo.launch()