test 003

app.py

@@ -310,50 +310,138 @@ def calculate_ctr(landmarks, corrected_landmarks=None):
     return round(ctr, 3), abs(tilt_angle)
 
 
+def detect_image_rotation(img):
+    """Detect rotation angle of chest X-ray using basic image analysis"""
+    # Apply edge detection
+    edges = cv2.Canny((img * 255).astype(np.uint8), 50, 150)
+    
+    # Find lines using Hough transform
+    lines = cv2.HoughLines(edges, 1, np.pi/180, threshold=100)
+    
+    if lines is not None:
+        angles = []
+        for rho, theta in lines[:min(10, len(lines))]:  # Consider top 10 lines
+            angle = np.degrees(theta) - 90  # Convert to rotation angle
+            # Filter for nearly horizontal or vertical lines
+            if abs(angle) < 30 or abs(angle) > 60:
+                angles.append(angle)
+        
+        if angles:
+            # Take median angle to avoid outliers
+            rotation_angle = np.median(angles)
+            if abs(rotation_angle) > 2:  # Only if significant rotation
+                return rotation_angle
+    
+    return 0
+
+def rotate_image(img, angle):
+    """Rotate image by given angle"""
+    if abs(angle) < 1:
+        return img, 0
+    
+    h, w = img.shape[:2]
+    center = (w // 2, h // 2)
+    
+    # Get rotation matrix
+    rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
+    
+    # Calculate new dimensions
+    cos_angle = abs(rotation_matrix[0, 0])
+    sin_angle = abs(rotation_matrix[0, 1])
+    new_w = int((h * sin_angle) + (w * cos_angle))
+    new_h = int((h * cos_angle) + (w * sin_angle))
+    
+    # Adjust translation
+    rotation_matrix[0, 2] += (new_w / 2) - center[0]
+    rotation_matrix[1, 2] += (new_h / 2) - center[1]
+    
+    # Rotate image
+    rotated = cv2.warpAffine(img, rotation_matrix, (new_w, new_h), 
+                            borderMode=cv2.BORDER_CONSTANT, borderValue=0)
+    
+    return rotated, angle
+
 def segment(input_img):
     global hybrid, device
 
     if hybrid is None:
         hybrid = loadModel(device)
 
-    input_img = cv2.imread(input_img, 0) / 255.0
-    original_shape = input_img.shape[:2]
+    original_img = cv2.imread(input_img, 0) / 255.0
+    original_shape = original_img.shape[:2]
+    
+    # Step 1: Detect and correct rotation BEFORE AI processing
+    detected_rotation = detect_image_rotation(original_img)
+    
+    if abs(detected_rotation) > 2:
+        # Rotate image to make it upright for AI processing
+        corrected_img, rotation_applied = rotate_image(original_img, -detected_rotation)
+        processing_img = corrected_img
+        was_rotated = True
+    else:
+        processing_img = original_img
+        rotation_applied = 0
+        was_rotated = False
 
-    img, (h, w, padding) = preprocess(input_img)
+    # Step 2: Preprocess the (potentially corrected) image
+    img, (h, w, padding) = preprocess(processing_img)
 
+    # Step 3: AI segmentation on corrected image
     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
     output = removePreprocess(output, (h, w, padding))
+    
+    # Step 5: If we rotated the image, rotate landmarks back to original orientation
+    if was_rotated:
+        corrected_h, corrected_w = processing_img.shape[:2]
+        corrected_center = np.array([corrected_w/2, corrected_h/2])
+        output_rotated = rotate_points(output.astype(float), detected_rotation, corrected_center)
+        
+        # Adjust coordinates back to original image size and position
+        # This is a simplified approach - you might need more sophisticated coordinate transformation
+        scale_x = original_shape[1] / corrected_w
+        scale_y = original_shape[0] / corrected_h
+        output_rotated[:, 0] *= scale_x
+        output_rotated[:, 1] *= scale_y
+        
+        output = output_rotated.astype('int')
+    else:
+        output = output.astype('int')
 
-    output = output.astype('int')
-
-    outseg, corrected_data = drawOnTop(input_img, output, original_shape)
+    # Step 6: Draw results on original image
+    outseg, corrected_data = drawOnTop(original_img, output, original_shape)
 
     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 tilt warning to interpretation
+    # Add rotation info to interpretation
+    rotation_warning = ""
+    if was_rotated:
+        rotation_warning = f" (🔄 Image was rotated {detected_rotation:.1f}° for AI processing)"
+    
+    # Add remaining tilt warning (after AI processing correction)
     tilt_warning = ""
     if tilt_angle > 5:
-        tilt_warning = f" (⚠️ Image tilted {tilt_angle:.1f}° - measurement corrected)"
+        tilt_warning = f" (⚠️ Remaining tilt: {tilt_angle:.1f}°)"
     elif tilt_angle > 2:
-        tilt_warning = f" (Image tilted {tilt_angle:.1f}° - corrected)"
+        tilt_warning = f" (Minor tilt: {tilt_angle:.1f}°)"
     
     if ctr_value < 0.5:
-        interpretation = f"Normal{tilt_warning}"
+        interpretation = f"Normal{rotation_warning}{tilt_warning}"
     elif 0.51 <= ctr_value <= 0.55:
-        interpretation = f"Mild Cardiomegaly (CTR 51-55%){tilt_warning}"
+        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%){tilt_warning}"
+        interpretation = f"Moderate Cardiomegaly (CTR 56-60%){rotation_warning}{tilt_warning}"
     elif ctr_value > 0.60:
-        interpretation = f"Severe Cardiomegaly (CTR > 60%){tilt_warning}"
+        interpretation = f"Severe Cardiomegaly (CTR > 60%){rotation_warning}{tilt_warning}"
     else:
-        interpretation = f"Cardiomegaly{tilt_warning}"
+        interpretation = f"Cardiomegaly{rotation_warning}{tilt_warning}"
 
     return outseg, "tmp/overlap_segmentation.png", ctr_value, interpretation