Update periodontitis_detection.py

periodontitis_detection.py

@@ -1,646 +1,646 @@
-import os
-import cv2
-import numpy as np
-import matplotlib.pyplot as plt
-import tensorflow as tf
-from ultralytics import YOLO
-
-class SimpleDentalSegmentationNoEnhance:
-    def __init__(self, unet_model_path, yolo_model_path, unet_input_size=(224,224,3)):
-        # Load TFLite U-Net
-        self.interpreter = tf.lite.Interpreter(model_path=unet_model_path)
-        self.interpreter.allocate_tensors()
-        self.input_details = self.interpreter.get_input_details()
-        self.output_details = self.interpreter.get_output_details()
-
-        # Force/prefer the desired U-Net input size
-        self.in_h, self.in_w, self.in_c = unet_input_size
-
-        # Load YOLOv8
-        self.yolo = YOLO(yolo_model_path)
-
-        print("Models loaded successfully!")
-        print(f"Using forced U-Net input shape: (1, {self.in_h}, {self.in_w}, {self.in_c})")
-        print(f"U-Net output shape (raw): {self.output_details[0]['shape']}")
-
-    def preprocess_for_unet(self, image_bgr):
-        img = image_bgr.copy()
-        proc_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-        proc_resized = cv2.resize(proc_rgb, (self.in_w, self.in_h), interpolation=cv2.INTER_LINEAR)
-        normalized = proc_resized.astype(np.float32) / 255.0
-        input_tensor = np.expand_dims(normalized, axis=0).astype(np.float32)
-        return input_tensor, proc_resized
-
-    def run_unet(self, image_bgr):
-        input_tensor, model_resized_image = self.preprocess_for_unet(image_bgr)
-
-        try:
-            self.interpreter.set_tensor(self.input_details[0]['index'], input_tensor)
-            self.interpreter.invoke()
-            output = self.interpreter.get_tensor(self.output_details[0]['index'])
-        except Exception as e:
-            print("Interpreter set_tensor failed, attempting to resize input to forced shape:", e)
-            try:
-                self.interpreter.resize_tensor_input(self.input_details[0]['index'], [1, self.in_h, self.in_w, self.in_c])
-                self.interpreter.allocate_tensors()
-                self.interpreter.set_tensor(self.input_details[0]['index'], input_tensor)
-                self.interpreter.invoke()
-                output = self.interpreter.get_tensor(self.output_details[0]['index'])
-            except Exception as e2:
-                raise RuntimeError("Failed to run TFLite interpreter") from e2
-
-        out = output[0]
-
-        if out.ndim == 3 and out.shape[2] >= 2:
-            class_map = np.argmax(out, axis=2).astype(np.uint8)
-            abc = (class_map == 1).astype(np.uint8)
-            cej = (class_map == 2).astype(np.uint8)
-        elif out.ndim == 2:
-            combined = out
-            abc = (combined > 0.5).astype(np.uint8)
-            cej = (combined > 0.8).astype(np.uint8)
-        else:
-            h_unet = out.shape[0]
-            w_unet = out.shape[1] if out.ndim >= 2 else (self.in_w)
-            abc = np.zeros((h_unet, w_unet), dtype=np.uint8)
-            cej = np.zeros((h_unet, w_unet), dtype=np.uint8)
-
-        return cej, abc, model_resized_image
-
-    def detect_teeth(self, image_bgr):
-        image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
-        results = self.yolo(image_rgb)
-        detections = []
-        for r in results:
-            boxes = getattr(r, "boxes", None)
-            if boxes is None:
-                continue
-            for i, box in enumerate(boxes):
-                try:
-                    xyxy = box.xyxy[0].cpu().numpy()
-                except Exception:
-                    xyxy = np.array(box.xyxy).astype(np.float32).reshape(-1)[:4]
-                try:
-                    conf = float(box.conf[0].cpu().numpy())
-                except Exception:
-                    conf = float(box.conf if hasattr(box, "conf") else 0.0)
-                detections.append({
-                    "bbox": xyxy.astype(np.float32),
-                    "confidence": conf,
-                    "tooth_id": len(detections) + 1
-                })
-        return detections
-
-    def resize_mask_to_original(self, mask, original_shape):
-        h_orig, w_orig = original_shape
-        mask_resized = cv2.resize((mask * 255).astype(np.uint8), (w_orig, h_orig), interpolation=cv2.INTER_NEAREST)
-        mask_resized = (mask_resized.astype(np.float32) / 255.0)
-        return mask_resized
-
-    def extract_abc_uppermost_line_within_bbox(self, abc_mask, bbox):
-        x1, y1, x2, y2 = bbox
-        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
-        height, width = abc_mask.shape
-        x1 = max(0, x1)
-        y1 = max(0, y1)
-        x2 = min(width-1, x2)
-        y2 = min(height-1, y2)
-        
-        abc_points = []
-        for x in range(x1, x2+1):
-            column_abc = np.where(abc_mask[y1:y2+1, x] == 255)[0]
-            if len(column_abc) > 0:
-                y_min_relative = np.min(column_abc)
-                y_absolute = y1 + y_min_relative
-                abc_points.append([x, y_absolute])
-        
-        if len(abc_points) < 2:
-            return None
-        return np.array(abc_points, dtype=np.int32).reshape(-1, 1, 2)
-
-    def extract_cej_lowermost_line_within_bbox(self, cej_mask, bbox):
-        x1, y1, x2, y2 = bbox
-        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
-        height, width = cej_mask.shape
-        x1 = max(0, x1)
-        y1 = max(0, y1)
-        x2 = min(width-1, x2)
-        y2 = min(height-1, y2)
-        
-        cej_points = []
-        for x in range(x1, x2+1):
-            column_cej = np.where(cej_mask[y1:y2+1, x] == 255)[0]
-            if len(column_cej) > 0:
-                y_max_relative = np.max(column_cej)
-                y_absolute = y1 + y_max_relative
-                cej_points.append([x, y_absolute])
-        
-        if len(cej_points) < 2:
-            return None
-        return np.array(cej_points, dtype=np.int32).reshape(-1, 1, 2)
-
-    def smooth_landmarks(self, points, window_size=5):
-        if points is None or len(points) < window_size:
-            return points
-        points_2d = points.reshape(-1, 2)
-        smoothed_points = []
-        for i in range(len(points_2d)):
-            start_idx = max(0, i - window_size // 2)
-            end_idx = min(len(points_2d), i + window_size // 2 + 1)
-            window_points = points_2d[start_idx:end_idx]
-            smoothed_y = np.mean(window_points[:, 1])
-            smoothed_points.append([points_2d[i][0], smoothed_y])
-        return np.array(smoothed_points, dtype=np.int32).reshape(-1, 1, 2)
-
-    def compute_cej_abc_distances(self, cej_points, abc_points):
-        """
-        Compute distances between CEJ and ABC points.
-        For each x-coordinate, find the vertical distance between CEJ and ABC.
-        """
-        if cej_points is None or abc_points is None:
-            return None
-            
-        cej_2d = cej_points.reshape(-1, 2)
-        abc_2d = abc_points.reshape(-1, 2)
-        
-        # Create dictionaries for quick lookup
-        cej_dict = {point[0]: point[1] for point in cej_2d}
-        abc_dict = {point[0]: point[1] for point in abc_2d}
-        
-        # Find common x-coordinates
-        common_x = set(cej_dict.keys()) & set(abc_dict.keys())
-        
-        if not common_x:
-            # If no exact matches, use interpolation
-            return self.compute_distances_with_interpolation(cej_2d, abc_2d)
-        
-        distances = []
-        connection_points = []
-        
-        for x in sorted(common_x):
-            cej_y = cej_dict[x]
-            abc_y = abc_dict[x]
-            distance = abs(abc_y - cej_y)  # Vertical distance
-            
-            distances.append({
-                'x': x,
-                'cej_y': cej_y,
-                'abc_y': abc_y,
-                'distance': distance
-            })
-            
-            connection_points.append([(x, cej_y), (x, abc_y)])
-        
-        return {
-            'distances': distances,
-            'connection_points': connection_points,
-            'mean_distance': np.mean([d['distance'] for d in distances]),
-            'max_distance': max([d['distance'] for d in distances]),
-            'min_distance': min([d['distance'] for d in distances])
-        }
-
-    def compute_distances_with_interpolation(self, cej_points, abc_points):
-        """
-        Compute distances using interpolation when points don't have exact x-matches.
-        """
-        # Get x-range that's common to both curves
-        cej_x_min, cej_x_max = np.min(cej_points[:, 0]), np.max(cej_points[:, 0])
-        abc_x_min, abc_x_max = np.min(abc_points[:, 0]), np.max(abc_points[:, 0])
-        
-        x_min = max(cej_x_min, abc_x_min)
-        x_max = min(cej_x_max, abc_x_max)
-        
-        if x_min >= x_max:
-            return None
-            
-        # Sample points at regular intervals
-        num_samples = min(50, x_max - x_min + 1)
-        x_sample = np.linspace(x_min, x_max, num_samples, dtype=int)
-        
-        # Interpolate y-values for both curves
-        cej_y_interp = np.interp(x_sample, cej_points[:, 0], cej_points[:, 1])
-        abc_y_interp = np.interp(x_sample, abc_points[:, 0], abc_points[:, 1])
-        
-        distances = []
-        connection_points = []
-        
-        for i, x in enumerate(x_sample):
-            cej_y = cej_y_interp[i]
-            abc_y = abc_y_interp[i]
-            distance = abs(abc_y - cej_y)
-            
-            distances.append({
-                'x': int(x),
-                'cej_y': int(cej_y),
-                'abc_y': int(abc_y),
-                'distance': distance
-            })
-            
-            connection_points.append([(int(x), int(cej_y)), (int(x), int(abc_y))])
-        
-        return {
-            'distances': distances,
-            'connection_points': connection_points,
-            'mean_distance': np.mean([d['distance'] for d in distances]),
-            'max_distance': max([d['distance'] for d in distances]),
-            'min_distance': min([d['distance'] for d in distances])
-        }
-
-    def draw_distance_measurements(self, image, distance_analysis, tooth_id):
-        """
-        Draw clean distance measurements on the image without text overlays.
-        """
-        if distance_analysis is None:
-            return image
-            
-        img_with_distances = image.copy()
-        
-        # Draw connection lines with gradient effect
-        connection_points = distance_analysis['connection_points']
-        distances = [d['distance'] for d in distance_analysis['distances']]
-        
-        if not distances:
-            return img_with_distances
-            
-        # Normalize distances for color mapping
-        min_dist = min(distances)
-        max_dist = max(distances)
-        dist_range = max_dist - min_dist if max_dist != min_dist else 1
-        
-        # Draw every 3rd line to reduce clutter, with color coding
-        for i in range(0, len(connection_points), 3):
-            start_point, end_point = connection_points[i]
-            distance = distances[i]
-            
-            # Color based on distance (green = small, red = large)
-            normalized_dist = (distance - min_dist) / dist_range
-            color_intensity = int(255 * normalized_dist)
-            color = (0, 255 - color_intensity, color_intensity)  # Green to Red
-            
-            # Draw thicker line for longer distances
-            thickness = max(1, int(3 * normalized_dist) + 1)
-            cv2.line(img_with_distances, start_point, end_point, color, thickness)
-            
-            # Add small circles at measurement points
-            cv2.circle(img_with_distances, start_point, 2, (255, 255, 255), -1)
-            cv2.circle(img_with_distances, end_point, 2, (255, 255, 255), -1)
-        
-        return img_with_distances
-
-    def analyze_image(self, image_path):
-        img_bgr = cv2.imread(image_path)
-        if img_bgr is None:
-            raise FileNotFoundError(f"Could not read image: {image_path}")
-        h_orig, w_orig = img_bgr.shape[:2]
-
-        cej_unet, abc_unet, _ = self.run_unet(img_bgr)
-
-        cej_orig = self.resize_mask_to_original(cej_unet, (h_orig, w_orig))
-        abc_orig = self.resize_mask_to_original(abc_unet, (h_orig, w_orig))
-
-        cej_bin = (cej_orig > 0.5).astype(np.uint8) * 255
-        abc_bin = (abc_orig > 0.5).astype(np.uint8) * 255
-
-        detections = self.detect_teeth(img_bgr)
-
-        combined = img_bgr.copy()
-        all_abc_segments = []
-        all_cej_segments = []
-        all_distance_analyses = []
-
-        for i, det in enumerate(detections):
-            x1, y1, x2, y2 = det["bbox"]
-            x1i = max(0, int(np.floor(x1)))
-            y1i = max(0, int(np.floor(y1)))
-            x2i = min(w_orig - 1, int(np.ceil(x2)))
-            y2i = min(h_orig - 1, int(np.ceil(y2)))
-            if x2i <= x1i or y2i <= y1i:
-                continue
-
-            cv2.rectangle(combined, (x1i, y1i), (x2i, y2i), (0,255,0), 2)
-
-            # ABC
-            abc_line_segment = self.extract_abc_uppermost_line_within_bbox(abc_bin, (x1i, y1i, x2i, y2i))
-            abc_data = None
-            if abc_line_segment is not None and len(abc_line_segment) > 1:
-                abc_line_segment = self.smooth_landmarks(abc_line_segment, window_size=3)
-                cv2.polylines(combined, [abc_line_segment], False, (255,0,0), 3)
-                abc_start = tuple(abc_line_segment[0][0])
-                abc_end = tuple(abc_line_segment[-1][0])
-                cv2.circle(combined, abc_start, 4, (255,165,0), -1)
-                cv2.circle(combined, abc_end, 4, (255,165,0), -1)
-                abc_data = {
-                    "points": abc_line_segment,
-                    "start": abc_start,
-                    "end": abc_end
-                }
-                all_abc_segments.append(abc_data)
-
-            # CEJ
-            cej_line_segment = self.extract_cej_lowermost_line_within_bbox(cej_bin, (x1i, y1i, x2i, y2i))
-            cej_data = None
-            if cej_line_segment is not None and len(cej_line_segment) > 1:
-                cej_line_segment = self.smooth_landmarks(cej_line_segment, window_size=3)
-                cv2.polylines(combined, [cej_line_segment], False, (0,0,255), 3)
-                cej_start = tuple(cej_line_segment[0][0])
-                cej_end = tuple(cej_line_segment[-1][0])
-                cv2.circle(combined, cej_start, 4, (0,255,255), -1)
-                cv2.circle(combined, cej_end, 4, (0,255,255), -1)
-                cej_data = {
-                    "points": cej_line_segment,
-                    "start": cej_start,
-                    "end": cej_end
-                }
-                all_cej_segments.append(cej_data)
-
-            # Compute distances between CEJ and ABC
-            distance_analysis = None
-            if abc_data is not None and cej_data is not None:
-                distance_analysis = self.compute_cej_abc_distances(
-                    cej_data["points"], abc_data["points"]
-                )
-                if distance_analysis is not None:
-                    combined = self.draw_distance_measurements(combined, distance_analysis, i+1)
-                    
-            all_distance_analyses.append({
-                'tooth_id': i + 1,
-                'analysis': distance_analysis
-            })
-
-            cv2.putText(combined, f"T{i+1}", (x1i+5, y1i+20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
-
-        results = {
-            "original": img_bgr,
-            "cej_mask": cej_bin,
-            "abc_mask": abc_bin,
-            "detections": detections,
-            "combined": combined,
-            "abc_segments": all_abc_segments,
-            "cej_segments": all_cej_segments,
-            "distance_analyses": all_distance_analyses
-        }
-        return results
-
-    def print_distance_summary(self, results):
-        """
-        Print a summary of distance measurements for all teeth.
-        """
-        print("\n" + "="*50)
-        print("CEJ-ABC DISTANCE ANALYSIS SUMMARY")
-        print("="*50)
-        
-        for tooth_data in results["distance_analyses"]:
-            tooth_id = tooth_data['tooth_id']
-            analysis = tooth_data['analysis']
-            
-            if analysis is None:
-                print(f"Tooth {tooth_id}: No distance measurements available")
-                continue
-                
-            print(f"\nTooth {tooth_id}:")
-            print(f"  Mean distance: {analysis['mean_distance']:.2f} pixels")
-            print(f"  Maximum distance: {analysis['max_distance']:.2f} pixels")
-            print(f"  Minimum distance: {analysis['min_distance']:.2f} pixels")
-            print(f"  Number of measurement points: {len(analysis['distances'])}")
-            
-            # Show some sample measurements
-            if len(analysis['distances']) > 0:
-                sample_size = min(3, len(analysis['distances']))
-                print(f"  Sample measurements:")
-                for i in range(0, len(analysis['distances']), len(analysis['distances'])//sample_size):
-                    d = analysis['distances'][i]
-                    print(f"    X={d['x']}: CEJ_Y={d['cej_y']}, ABC_Y={d['abc_y']}, Distance={d['distance']:.1f}px")
-
-    def create_distance_heatmap(self, results):
-        """Create a heatmap visualization of distances across all teeth."""
-        all_distances = []
-        tooth_labels = []
-        
-        for tooth_data in results["distance_analyses"]:
-            if tooth_data['analysis'] is not None:
-                distances = [d['distance'] for d in tooth_data['analysis']['distances']]
-                all_distances.extend(distances)
-                tooth_labels.extend([f"T{tooth_data['tooth_id']}"] * len(distances))
-        
-        if not all_distances:
-            return None
-            
-        # Create histogram data
-        unique_teeth = list(set(tooth_labels))
-        tooth_distances = {tooth: [] for tooth in unique_teeth}
-        
-        for i, tooth in enumerate(tooth_labels):
-            tooth_distances[tooth].append(all_distances[i])
-        
-        return tooth_distances
-
-    def create_overlay_image(self, results):
-        """Create an enhanced overlay image with better visualization."""
-        img = results["original"].copy()
-        cej_mask = results["cej_mask"]
-        abc_mask = results["abc_mask"]
-        
-        # Create colored overlays
-        overlay = img.copy()
-        
-        # CEJ in red with transparency
-        cej_colored = np.zeros_like(img)
-        cej_colored[:, :, 2] = cej_mask  # Red channel
-        
-        # ABC in blue with transparency
-        abc_colored = np.zeros_like(img)
-        abc_colored[:, :, 0] = abc_mask  # Blue channel
-        
-        # Blend overlays
-        alpha = 0.4
-        overlay = cv2.addWeighted(overlay, 1-alpha, cej_colored, alpha, 0)
-        overlay = cv2.addWeighted(overlay, 1-alpha, abc_colored, alpha, 0)
-        
-        # Add tooth detection boxes and labels
-        for i, det in enumerate(results["detections"]):
-            x1, y1, x2, y2 = det["bbox"].astype(int)
-            cv2.rectangle(overlay, (x1, y1), (x2, y2), (0, 255, 0), 2)
-            
-            # Add tooth label with background
-            label = f"Tooth {i+1}"
-            font = cv2.FONT_HERSHEY_SIMPLEX
-            font_scale = 0.7
-            thickness = 2
-            (text_width, text_height), _ = cv2.getTextSize(label, font, font_scale, thickness)
-            
-            cv2.rectangle(overlay, (x1, y1-text_height-10), (x1+text_width+10, y1), (0, 255, 0), -1)
-            cv2.putText(overlay, label, (x1+5, y1-5), font, font_scale, (0, 0, 0), thickness)
-        
-        return overlay
-
-    def visualize_results(self, results, save_path=None):
-        # Prepare images
-        orig_rgb = cv2.cvtColor(results["original"], cv2.COLOR_BGR2RGB)
-        combined_rgb = cv2.cvtColor(results["combined"], cv2.COLOR_BGR2RGB)
-        overlay_bgr = self.create_overlay_image(results)
-        overlay_rgb = cv2.cvtColor(overlay_bgr, cv2.COLOR_BGR2RGB)
-        
-        # Create figure with custom layout
-        fig = plt.figure(figsize=(20, 15))
-        
-        # Create a grid layout
-        gs = fig.add_gridspec(3, 4, height_ratios=[1, 1, 0.8], hspace=0.3, wspace=0.2)
-        
-        # Top row - Original images
-        ax1 = fig.add_subplot(gs[0, 0])
-        ax1.imshow(orig_rgb)
-        ax1.set_title("Original Image", fontsize=14, fontweight='bold')
-        ax1.axis("off")
-        
-        ax2 = fig.add_subplot(gs[0, 1])
-        ax2.imshow(overlay_rgb)
-        ax2.set_title("Segmentation Overlay\n(Red: CEJ, Blue: ABC)", fontsize=14, fontweight='bold')
-        ax2.axis("off")
-        
-        ax3 = fig.add_subplot(gs[0, 2])
-        ax3.imshow(combined_rgb)
-        ax3.set_title("Distance Analysis\n(Yellow: Measurements)", fontsize=14, fontweight='bold')
-        ax3.axis("off")
-        
-        # Individual mask visualization
-        ax4 = fig.add_subplot(gs[0, 3])
-        # Create combined mask visualization
-        combined_mask = np.zeros((*results["cej_mask"].shape, 3), dtype=np.uint8)
-        combined_mask[:, :, 2] = results["cej_mask"]  # CEJ in red
-        combined_mask[:, :, 0] = results["abc_mask"]  # ABC in blue
-        ax4.imshow(combined_mask)
-        ax4.set_title("Combined Masks\n(Red: CEJ, Blue: ABC)", fontsize=14, fontweight='bold')
-        ax4.axis("off")
-        
-        # Middle row - Distance analysis charts
-        ax5 = fig.add_subplot(gs[1, :2])  # Span 2 columns
-        
-        # Create bar chart of average distances per tooth
-        tooth_means = []
-        tooth_labels = []
-        tooth_maxs = []
-        tooth_mins = []
-        
-        for tooth_data in results["distance_analyses"]:
-            if tooth_data['analysis'] is not None:
-                tooth_labels.append(f"T{tooth_data['tooth_id']}")
-                tooth_means.append(tooth_data['analysis']['mean_distance'])
-                tooth_maxs.append(tooth_data['analysis']['max_distance'])
-                tooth_mins.append(tooth_data['analysis']['min_distance'])
-        
-        if tooth_means:
-            x_pos = np.arange(len(tooth_labels))
-            bars = ax5.bar(x_pos, tooth_means, alpha=0.7, color='skyblue', edgecolor='navy')
-            
-            # Add error bars showing min/max range
-            yerr_lower = np.array(tooth_means) - np.array(tooth_mins)
-            yerr_upper = np.array(tooth_maxs) - np.array(tooth_means)
-            ax5.errorbar(x_pos, tooth_means, yerr=[yerr_lower, yerr_upper], 
-                        fmt='none', ecolor='red', capsize=5, alpha=0.7)
-            
-            ax5.set_xlabel("Tooth Number", fontsize=12, fontweight='bold')
-            ax5.set_ylabel("Distance (pixels)", fontsize=12, fontweight='bold')
-            ax5.set_title("CEJ-ABC Distance Analysis by Tooth", fontsize=14, fontweight='bold')
-            ax5.set_xticks(x_pos)
-            ax5.set_xticklabels(tooth_labels)
-            ax5.grid(True, alpha=0.3)
-            
-            # Add value labels on bars
-            for i, (bar, mean_val, max_val, min_val) in enumerate(zip(bars, tooth_means, tooth_maxs, tooth_mins)):
-                height = bar.get_height()
-                ax5.text(bar.get_x() + bar.get_width()/2., height + 1,
-                        f'{mean_val:.1f}', ha='center', va='bottom', fontweight='bold')
-        else:
-            ax5.text(0.5, 0.5, "No distance measurements available", 
-                    ha='center', va='center', transform=ax5.transAxes, fontsize=14)
-            ax5.set_title("CEJ-ABC Distance Analysis", fontsize=14, fontweight='bold')
-        
-        # Distance distribution histogram
-        ax6 = fig.add_subplot(gs[1, 2:])  # Span 2 columns
-        
-        tooth_distances = self.create_distance_heatmap(results)
-        if tooth_distances:
-            all_vals = []
-            labels = []
-            colors = plt.cm.Set3(np.linspace(0, 1, len(tooth_distances)))
-            
-            for i, (tooth, distances) in enumerate(tooth_distances.items()):
-                ax6.hist(distances, bins=20, alpha=0.6, label=tooth, color=colors[i])
-                all_vals.extend(distances)
-            
-            ax6.set_xlabel("Distance (pixels)", fontsize=12, fontweight='bold')
-            ax6.set_ylabel("Frequency", fontsize=12, fontweight='bold')
-            ax6.set_title("Distance Distribution Across All Measurements", fontsize=14, fontweight='bold')
-            ax6.legend()
-            ax6.grid(True, alpha=0.3)
-        else:
-            ax6.text(0.5, 0.5, "No distance data available for histogram", 
-                    ha='center', va='center', transform=ax6.transAxes, fontsize=14)
-            ax6.set_title("Distance Distribution", fontsize=14, fontweight='bold')
-        
-        # Bottom row - Summary statistics table
-        ax7 = fig.add_subplot(gs[2, :])
-        ax7.axis('tight')
-        ax7.axis('off')
-        
-        # Create summary table
-        if tooth_means:
-            table_data = []
-            headers = ['Tooth', 'Mean Distance (px)', 'Max Distance (px)', 'Min Distance (px)', 'Range (px)', 'Measurements']
-            
-            for tooth_data in results["distance_analyses"]:
-                if tooth_data['analysis'] is not None:
-                    analysis = tooth_data['analysis']
-                    range_val = analysis['max_distance'] - analysis['min_distance']
-                    num_measurements = len(analysis['distances'])
-                    
-                    table_data.append([
-                        f"T{tooth_data['tooth_id']}",
-                        f"{analysis['mean_distance']:.2f}",
-                        f"{analysis['max_distance']:.2f}",
-                        f"{analysis['min_distance']:.2f}",
-                        f"{range_val:.2f}",
-                        str(num_measurements)
-                    ])
-            
-            table = ax7.table(cellText=table_data, colLabels=headers, 
-                             cellLoc='center', loc='center')
-            table.auto_set_font_size(False)
-            table.set_fontsize(10)
-            table.scale(1, 2)
-            
-            # Style the table
-            for (i, j), cell in table.get_celld().items():
-                if i == 0:  # Header row
-                    cell.set_text_props(weight='bold', color='white')
-                    cell.set_facecolor('#4472C4')
-                else:
-                    cell.set_facecolor('#F2F2F2' if i % 2 == 0 else 'white')
-                    
-            ax7.set_title("Detailed Distance Analysis Summary", fontsize=14, fontweight='bold', pad=20)
-        else:
-            ax7.text(0.5, 0.5, "No measurements available for summary table", 
-                    ha='center', va='center', transform=ax7.transAxes, fontsize=14)
-        
-        plt.suptitle("Comprehensive Dental CEJ-ABC Distance Analysis", fontsize=16, fontweight='bold', y=0.98)
-        
-        if save_path:
-            plt.savefig(save_path, dpi=300, bbox_inches="tight", facecolor='white')
-            print(f"Saved enhanced visualization to {save_path}")
-        
-        return fig
-
-
-if __name__ == "__main__":
-    unet_model = "models/dental_segmentation_model_augment_100epochs.tflite"
-    yolo_model = "models/best v8n-seg_float16.tflite"
-    image_path = "trial2.jpg"
-
-    seg = SimpleDentalSegmentationNoEnhance(unet_model, yolo_model)
-    res = seg.analyze_image(image_path)
-    
-    # Print distance analysis summary
-    seg.print_distance_summary(res)
-    
-    fig = seg.visualize_results(res, save_path="segmentation_with_distances.png")
+import os
+import cv2
+import numpy as np
+import matplotlib.pyplot as plt
+import tensorflow as tf
+from ultralytics import YOLO
+
+class SimpleDentalSegmentationNoEnhance:
+    def __init__(self, unet_model_path, yolo_model_path, unet_input_size=(224,224,3)):
+        # Load TFLite U-Net
+        self.interpreter = tf.lite.Interpreter(model_path=unet_model_path)
+        self.interpreter.allocate_tensors()
+        self.input_details = self.interpreter.get_input_details()
+        self.output_details = self.interpreter.get_output_details()
+
+        # Force/prefer the desired U-Net input size
+        self.in_h, self.in_w, self.in_c = unet_input_size
+
+        # Load YOLOv8
+        self.yolo = YOLO(yolo_model_path)
+
+        print("Models loaded successfully!")
+        print(f"Using forced U-Net input shape: (1, {self.in_h}, {self.in_w}, {self.in_c})")
+        print(f"U-Net output shape (raw): {self.output_details[0]['shape']}")
+
+    def preprocess_for_unet(self, image_bgr):
+        img = image_bgr.copy()
+        proc_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        proc_resized = cv2.resize(proc_rgb, (self.in_w, self.in_h), interpolation=cv2.INTER_LINEAR)
+        normalized = proc_resized.astype(np.float32) / 255.0
+        input_tensor = np.expand_dims(normalized, axis=0).astype(np.float32)
+        return input_tensor, proc_resized
+
+    def run_unet(self, image_bgr):
+        input_tensor, model_resized_image = self.preprocess_for_unet(image_bgr)
+
+        try:
+            self.interpreter.set_tensor(self.input_details[0]['index'], input_tensor)
+            self.interpreter.invoke()
+            output = self.interpreter.get_tensor(self.output_details[0]['index'])
+        except Exception as e:
+            print("Interpreter set_tensor failed, attempting to resize input to forced shape:", e)
+            try:
+                self.interpreter.resize_tensor_input(self.input_details[0]['index'], [1, self.in_h, self.in_w, self.in_c])
+                self.interpreter.allocate_tensors()
+                self.interpreter.set_tensor(self.input_details[0]['index'], input_tensor)
+                self.interpreter.invoke()
+                output = self.interpreter.get_tensor(self.output_details[0]['index'])
+            except Exception as e2:
+                raise RuntimeError("Failed to run TFLite interpreter") from e2
+
+        out = output[0]
+
+        if out.ndim == 3 and out.shape[2] >= 2:
+            class_map = np.argmax(out, axis=2).astype(np.uint8)
+            abc = (class_map == 1).astype(np.uint8)
+            cej = (class_map == 2).astype(np.uint8)
+        elif out.ndim == 2:
+            combined = out
+            abc = (combined > 0.5).astype(np.uint8)
+            cej = (combined > 0.8).astype(np.uint8)
+        else:
+            h_unet = out.shape[0]
+            w_unet = out.shape[1] if out.ndim >= 2 else (self.in_w)
+            abc = np.zeros((h_unet, w_unet), dtype=np.uint8)
+            cej = np.zeros((h_unet, w_unet), dtype=np.uint8)
+
+        return cej, abc, model_resized_image
+
+    def detect_teeth(self, image_bgr):
+        image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+        results = self.yolo(image_rgb)
+        detections = []
+        for r in results:
+            boxes = getattr(r, "boxes", None)
+            if boxes is None:
+                continue
+            for i, box in enumerate(boxes):
+                try:
+                    xyxy = box.xyxy[0].cpu().numpy()
+                except Exception:
+                    xyxy = np.array(box.xyxy).astype(np.float32).reshape(-1)[:4]
+                try:
+                    conf = float(box.conf[0].cpu().numpy())
+                except Exception:
+                    conf = float(box.conf if hasattr(box, "conf") else 0.0)
+                detections.append({
+                    "bbox": xyxy.astype(np.float32),
+                    "confidence": conf,
+                    "tooth_id": len(detections) + 1
+                })
+        return detections
+
+    def resize_mask_to_original(self, mask, original_shape):
+        h_orig, w_orig = original_shape
+        mask_resized = cv2.resize((mask * 255).astype(np.uint8), (w_orig, h_orig), interpolation=cv2.INTER_NEAREST)
+        mask_resized = (mask_resized.astype(np.float32) / 255.0)
+        return mask_resized
+
+    def extract_abc_uppermost_line_within_bbox(self, abc_mask, bbox):
+        x1, y1, x2, y2 = bbox
+        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+        height, width = abc_mask.shape
+        x1 = max(0, x1)
+        y1 = max(0, y1)
+        x2 = min(width-1, x2)
+        y2 = min(height-1, y2)
+        
+        abc_points = []
+        for x in range(x1, x2+1):
+            column_abc = np.where(abc_mask[y1:y2+1, x] == 255)[0]
+            if len(column_abc) > 0:
+                y_min_relative = np.min(column_abc)
+                y_absolute = y1 + y_min_relative
+                abc_points.append([x, y_absolute])
+        
+        if len(abc_points) < 2:
+            return None
+        return np.array(abc_points, dtype=np.int32).reshape(-1, 1, 2)
+
+    def extract_cej_lowermost_line_within_bbox(self, cej_mask, bbox):
+        x1, y1, x2, y2 = bbox
+        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+        height, width = cej_mask.shape
+        x1 = max(0, x1)
+        y1 = max(0, y1)
+        x2 = min(width-1, x2)
+        y2 = min(height-1, y2)
+        
+        cej_points = []
+        for x in range(x1, x2+1):
+            column_cej = np.where(cej_mask[y1:y2+1, x] == 255)[0]
+            if len(column_cej) > 0:
+                y_max_relative = np.max(column_cej)
+                y_absolute = y1 + y_max_relative
+                cej_points.append([x, y_absolute])
+        
+        if len(cej_points) < 2:
+            return None
+        return np.array(cej_points, dtype=np.int32).reshape(-1, 1, 2)
+
+    def smooth_landmarks(self, points, window_size=5):
+        if points is None or len(points) < window_size:
+            return points
+        points_2d = points.reshape(-1, 2)
+        smoothed_points = []
+        for i in range(len(points_2d)):
+            start_idx = max(0, i - window_size // 2)
+            end_idx = min(len(points_2d), i + window_size // 2 + 1)
+            window_points = points_2d[start_idx:end_idx]
+            smoothed_y = np.mean(window_points[:, 1])
+            smoothed_points.append([points_2d[i][0], smoothed_y])
+        return np.array(smoothed_points, dtype=np.int32).reshape(-1, 1, 2)
+
+    def compute_cej_abc_distances(self, cej_points, abc_points):
+        """
+        Compute distances between CEJ and ABC points.
+        For each x-coordinate, find the vertical distance between CEJ and ABC.
+        """
+        if cej_points is None or abc_points is None:
+            return None
+            
+        cej_2d = cej_points.reshape(-1, 2)
+        abc_2d = abc_points.reshape(-1, 2)
+        
+        # Create dictionaries for quick lookup
+        cej_dict = {point[0]: point[1] for point in cej_2d}
+        abc_dict = {point[0]: point[1] for point in abc_2d}
+        
+        # Find common x-coordinates
+        common_x = set(cej_dict.keys()) & set(abc_dict.keys())
+        
+        if not common_x:
+            # If no exact matches, use interpolation
+            return self.compute_distances_with_interpolation(cej_2d, abc_2d)
+        
+        distances = []
+        connection_points = []
+        
+        for x in sorted(common_x):
+            cej_y = cej_dict[x]
+            abc_y = abc_dict[x]
+            distance = abs(abc_y - cej_y)  # Vertical distance
+            
+            distances.append({
+                'x': x,
+                'cej_y': cej_y,
+                'abc_y': abc_y,
+                'distance': distance
+            })
+            
+            connection_points.append([(x, cej_y), (x, abc_y)])
+        
+        return {
+            'distances': distances,
+            'connection_points': connection_points,
+            'mean_distance': np.mean([d['distance'] for d in distances]),
+            'max_distance': max([d['distance'] for d in distances]),
+            'min_distance': min([d['distance'] for d in distances])
+        }
+
+    def compute_distances_with_interpolation(self, cej_points, abc_points):
+        """
+        Compute distances using interpolation when points don't have exact x-matches.
+        """
+        # Get x-range that's common to both curves
+        cej_x_min, cej_x_max = np.min(cej_points[:, 0]), np.max(cej_points[:, 0])
+        abc_x_min, abc_x_max = np.min(abc_points[:, 0]), np.max(abc_points[:, 0])
+        
+        x_min = max(cej_x_min, abc_x_min)
+        x_max = min(cej_x_max, abc_x_max)
+        
+        if x_min >= x_max:
+            return None
+            
+        # Sample points at regular intervals
+        num_samples = min(50, x_max - x_min + 1)
+        x_sample = np.linspace(x_min, x_max, num_samples, dtype=int)
+        
+        # Interpolate y-values for both curves
+        cej_y_interp = np.interp(x_sample, cej_points[:, 0], cej_points[:, 1])
+        abc_y_interp = np.interp(x_sample, abc_points[:, 0], abc_points[:, 1])
+        
+        distances = []
+        connection_points = []
+        
+        for i, x in enumerate(x_sample):
+            cej_y = cej_y_interp[i]
+            abc_y = abc_y_interp[i]
+            distance = abs(abc_y - cej_y)
+            
+            distances.append({
+                'x': int(x),
+                'cej_y': int(cej_y),
+                'abc_y': int(abc_y),
+                'distance': distance
+            })
+            
+            connection_points.append([(int(x), int(cej_y)), (int(x), int(abc_y))])
+        
+        return {
+            'distances': distances,
+            'connection_points': connection_points,
+            'mean_distance': np.mean([d['distance'] for d in distances]),
+            'max_distance': max([d['distance'] for d in distances]),
+            'min_distance': min([d['distance'] for d in distances])
+        }
+
+    def draw_distance_measurements(self, image, distance_analysis, tooth_id):
+        """
+        Draw clean distance measurements on the image without text overlays.
+        """
+        if distance_analysis is None:
+            return image
+            
+        img_with_distances = image.copy()
+        
+        # Draw connection lines with gradient effect
+        connection_points = distance_analysis['connection_points']
+        distances = [d['distance'] for d in distance_analysis['distances']]
+        
+        if not distances:
+            return img_with_distances
+            
+        # Normalize distances for color mapping
+        min_dist = min(distances)
+        max_dist = max(distances)
+        dist_range = max_dist - min_dist if max_dist != min_dist else 1
+        
+        # Draw every 3rd line to reduce clutter, with color coding
+        for i in range(0, len(connection_points), 3):
+            start_point, end_point = connection_points[i]
+            distance = distances[i]
+            
+            # Color based on distance (green = small, red = large)
+            normalized_dist = (distance - min_dist) / dist_range
+            color_intensity = int(255 * normalized_dist)
+            color = (0, 255 - color_intensity, color_intensity)  # Green to Red
+            
+            # Draw thicker line for longer distances
+            thickness = max(1, int(3 * normalized_dist) + 1)
+            cv2.line(img_with_distances, start_point, end_point, color, thickness)
+            
+            # Add small circles at measurement points
+            cv2.circle(img_with_distances, start_point, 2, (255, 255, 255), -1)
+            cv2.circle(img_with_distances, end_point, 2, (255, 255, 255), -1)
+        
+        return img_with_distances
+
+    def analyze_image(self, image_path):
+        img_bgr = cv2.imread(image_path)
+        if img_bgr is None:
+            raise FileNotFoundError(f"Could not read image: {image_path}")
+        h_orig, w_orig = img_bgr.shape[:2]
+
+        cej_unet, abc_unet, _ = self.run_unet(img_bgr)
+
+        cej_orig = self.resize_mask_to_original(cej_unet, (h_orig, w_orig))
+        abc_orig = self.resize_mask_to_original(abc_unet, (h_orig, w_orig))
+
+        cej_bin = (cej_orig > 0.5).astype(np.uint8) * 255
+        abc_bin = (abc_orig > 0.5).astype(np.uint8) * 255
+
+        detections = self.detect_teeth(img_bgr)
+
+        combined = img_bgr.copy()
+        all_abc_segments = []
+        all_cej_segments = []
+        all_distance_analyses = []
+
+        for i, det in enumerate(detections):
+            x1, y1, x2, y2 = det["bbox"]
+            x1i = max(0, int(np.floor(x1)))
+            y1i = max(0, int(np.floor(y1)))
+            x2i = min(w_orig - 1, int(np.ceil(x2)))
+            y2i = min(h_orig - 1, int(np.ceil(y2)))
+            if x2i <= x1i or y2i <= y1i:
+                continue
+
+            cv2.rectangle(combined, (x1i, y1i), (x2i, y2i), (0,255,0), 2)
+
+            # ABC
+            abc_line_segment = self.extract_abc_uppermost_line_within_bbox(abc_bin, (x1i, y1i, x2i, y2i))
+            abc_data = None
+            if abc_line_segment is not None and len(abc_line_segment) > 1:
+                abc_line_segment = self.smooth_landmarks(abc_line_segment, window_size=3)
+                cv2.polylines(combined, [abc_line_segment], False, (255,0,0), 3)
+                abc_start = tuple(abc_line_segment[0][0])
+                abc_end = tuple(abc_line_segment[-1][0])
+                cv2.circle(combined, abc_start, 4, (255,165,0), -1)
+                cv2.circle(combined, abc_end, 4, (255,165,0), -1)
+                abc_data = {
+                    "points": abc_line_segment,
+                    "start": abc_start,
+                    "end": abc_end
+                }
+                all_abc_segments.append(abc_data)
+
+            # CEJ
+            cej_line_segment = self.extract_cej_lowermost_line_within_bbox(cej_bin, (x1i, y1i, x2i, y2i))
+            cej_data = None
+            if cej_line_segment is not None and len(cej_line_segment) > 1:
+                cej_line_segment = self.smooth_landmarks(cej_line_segment, window_size=3)
+                cv2.polylines(combined, [cej_line_segment], False, (0,0,255), 3)
+                cej_start = tuple(cej_line_segment[0][0])
+                cej_end = tuple(cej_line_segment[-1][0])
+                cv2.circle(combined, cej_start, 4, (0,255,255), -1)
+                cv2.circle(combined, cej_end, 4, (0,255,255), -1)
+                cej_data = {
+                    "points": cej_line_segment,
+                    "start": cej_start,
+                    "end": cej_end
+                }
+                all_cej_segments.append(cej_data)
+
+            # Compute distances between CEJ and ABC
+            distance_analysis = None
+            if abc_data is not None and cej_data is not None:
+                distance_analysis = self.compute_cej_abc_distances(
+                    cej_data["points"], abc_data["points"]
+                )
+                if distance_analysis is not None:
+                    combined = self.draw_distance_measurements(combined, distance_analysis, i+1)
+                    
+            all_distance_analyses.append({
+                'tooth_id': i + 1,
+                'analysis': distance_analysis
+            })
+
+            cv2.putText(combined, f"T{i+1}", (x1i+5, y1i+20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
+
+        results = {
+            "original": img_bgr,
+            "cej_mask": cej_bin,
+            "abc_mask": abc_bin,
+            "detections": detections,
+            "combined": combined,
+            "abc_segments": all_abc_segments,
+            "cej_segments": all_cej_segments,
+            "distance_analyses": all_distance_analyses
+        }
+        return results
+
+    def print_distance_summary(self, results):
+        """
+        Print a summary of distance measurements for all teeth.
+        """
+        print("\n" + "="*50)
+        print("CEJ-ABC DISTANCE ANALYSIS SUMMARY")
+        print("="*50)
+        
+        for tooth_data in results["distance_analyses"]:
+            tooth_id = tooth_data['tooth_id']
+            analysis = tooth_data['analysis']
+            
+            if analysis is None:
+                print(f"Tooth {tooth_id}: No distance measurements available")
+                continue
+                
+            print(f"\nTooth {tooth_id}:")
+            print(f"  Mean distance: {analysis['mean_distance']:.2f} pixels")
+            print(f"  Maximum distance: {analysis['max_distance']:.2f} pixels")
+            print(f"  Minimum distance: {analysis['min_distance']:.2f} pixels")
+            print(f"  Number of measurement points: {len(analysis['distances'])}")
+            
+            # Show some sample measurements
+            if len(analysis['distances']) > 0:
+                sample_size = min(3, len(analysis['distances']))
+                print(f"  Sample measurements:")
+                for i in range(0, len(analysis['distances']), len(analysis['distances'])//sample_size):
+                    d = analysis['distances'][i]
+                    print(f"    X={d['x']}: CEJ_Y={d['cej_y']}, ABC_Y={d['abc_y']}, Distance={d['distance']:.1f}px")
+
+    def create_distance_heatmap(self, results):
+        """Create a heatmap visualization of distances across all teeth."""
+        all_distances = []
+        tooth_labels = []
+        
+        for tooth_data in results["distance_analyses"]:
+            if tooth_data['analysis'] is not None:
+                distances = [d['distance'] for d in tooth_data['analysis']['distances']]
+                all_distances.extend(distances)
+                tooth_labels.extend([f"T{tooth_data['tooth_id']}"] * len(distances))
+        
+        if not all_distances:
+            return None
+            
+        # Create histogram data
+        unique_teeth = list(set(tooth_labels))
+        tooth_distances = {tooth: [] for tooth in unique_teeth}
+        
+        for i, tooth in enumerate(tooth_labels):
+            tooth_distances[tooth].append(all_distances[i])
+        
+        return tooth_distances
+
+    def create_overlay_image(self, results):
+        """Create an enhanced overlay image with better visualization."""
+        img = results["original"].copy()
+        cej_mask = results["cej_mask"]
+        abc_mask = results["abc_mask"]
+        
+        # Create colored overlays
+        overlay = img.copy()
+        
+        # CEJ in red with transparency
+        cej_colored = np.zeros_like(img)
+        cej_colored[:, :, 2] = cej_mask  # Red channel
+        
+        # ABC in blue with transparency
+        abc_colored = np.zeros_like(img)
+        abc_colored[:, :, 0] = abc_mask  # Blue channel
+        
+        # Blend overlays
+        alpha = 0.4
+        overlay = cv2.addWeighted(overlay, 1-alpha, cej_colored, alpha, 0)
+        overlay = cv2.addWeighted(overlay, 1-alpha, abc_colored, alpha, 0)
+        
+        # Add tooth detection boxes and labels
+        for i, det in enumerate(results["detections"]):
+            x1, y1, x2, y2 = det["bbox"].astype(int)
+            cv2.rectangle(overlay, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            
+            # Add tooth label with background
+            label = f"Tooth {i+1}"
+            font = cv2.FONT_HERSHEY_SIMPLEX
+            font_scale = 0.7
+            thickness = 2
+            (text_width, text_height), _ = cv2.getTextSize(label, font, font_scale, thickness)
+            
+            cv2.rectangle(overlay, (x1, y1-text_height-10), (x1+text_width+10, y1), (0, 255, 0), -1)
+            cv2.putText(overlay, label, (x1+5, y1-5), font, font_scale, (0, 0, 0), thickness)
+        
+        return overlay
+
+    def visualize_results(self, results, save_path=None):
+        # Prepare images
+        orig_rgb = cv2.cvtColor(results["original"], cv2.COLOR_BGR2RGB)
+        combined_rgb = cv2.cvtColor(results["combined"], cv2.COLOR_BGR2RGB)
+        overlay_bgr = self.create_overlay_image(results)
+        overlay_rgb = cv2.cvtColor(overlay_bgr, cv2.COLOR_BGR2RGB)
+        
+        # Create figure with custom layout
+        fig = plt.figure(figsize=(20, 15))
+        
+        # Create a grid layout
+        gs = fig.add_gridspec(3, 4, height_ratios=[1, 1, 0.8], hspace=0.3, wspace=0.2)
+        
+        # Top row - Original images
+        ax1 = fig.add_subplot(gs[0, 0])
+        ax1.imshow(orig_rgb)
+        ax1.set_title("Original Image", fontsize=14, fontweight='bold')
+        ax1.axis("off")
+        
+        ax2 = fig.add_subplot(gs[0, 1])
+        ax2.imshow(overlay_rgb)
+        ax2.set_title("Segmentation Overlay\n(Red: CEJ, Blue: ABC)", fontsize=14, fontweight='bold')
+        ax2.axis("off")
+        
+        ax3 = fig.add_subplot(gs[0, 2])
+        ax3.imshow(combined_rgb)
+        ax3.set_title("Distance Analysis\n(Yellow: Measurements)", fontsize=14, fontweight='bold')
+        ax3.axis("off")
+        
+        # Individual mask visualization
+        ax4 = fig.add_subplot(gs[0, 3])
+        # Create combined mask visualization
+        combined_mask = np.zeros((*results["cej_mask"].shape, 3), dtype=np.uint8)
+        combined_mask[:, :, 2] = results["cej_mask"]  # CEJ in red
+        combined_mask[:, :, 0] = results["abc_mask"]  # ABC in blue
+        ax4.imshow(combined_mask)
+        ax4.set_title("Combined Masks\n(Red: CEJ, Blue: ABC)", fontsize=14, fontweight='bold')
+        ax4.axis("off")
+        
+        # Middle row - Distance analysis charts
+        ax5 = fig.add_subplot(gs[1, :2])  # Span 2 columns
+        
+        # Create bar chart of average distances per tooth
+        tooth_means = []
+        tooth_labels = []
+        tooth_maxs = []
+        tooth_mins = []
+        
+        for tooth_data in results["distance_analyses"]:
+            if tooth_data['analysis'] is not None:
+                tooth_labels.append(f"T{tooth_data['tooth_id']}")
+                tooth_means.append(tooth_data['analysis']['mean_distance'])
+                tooth_maxs.append(tooth_data['analysis']['max_distance'])
+                tooth_mins.append(tooth_data['analysis']['min_distance'])
+        
+        if tooth_means:
+            x_pos = np.arange(len(tooth_labels))
+            bars = ax5.bar(x_pos, tooth_means, alpha=0.7, color='skyblue', edgecolor='navy')
+            
+            # Add error bars showing min/max range
+            yerr_lower = np.array(tooth_means) - np.array(tooth_mins)
+            yerr_upper = np.array(tooth_maxs) - np.array(tooth_means)
+            ax5.errorbar(x_pos, tooth_means, yerr=[yerr_lower, yerr_upper], 
+                        fmt='none', ecolor='red', capsize=5, alpha=0.7)
+            
+            ax5.set_xlabel("Tooth Number", fontsize=12, fontweight='bold')
+            ax5.set_ylabel("Distance (pixels)", fontsize=12, fontweight='bold')
+            ax5.set_title("CEJ-ABC Distance Analysis by Tooth", fontsize=14, fontweight='bold')
+            ax5.set_xticks(x_pos)
+            ax5.set_xticklabels(tooth_labels)
+            ax5.grid(True, alpha=0.3)
+            
+            # Add value labels on bars
+            for i, (bar, mean_val, max_val, min_val) in enumerate(zip(bars, tooth_means, tooth_maxs, tooth_mins)):
+                height = bar.get_height()
+                ax5.text(bar.get_x() + bar.get_width()/2., height + 1,
+                        f'{mean_val:.1f}', ha='center', va='bottom', fontweight='bold')
+        else:
+            ax5.text(0.5, 0.5, "No distance measurements available", 
+                    ha='center', va='center', transform=ax5.transAxes, fontsize=14)
+            ax5.set_title("CEJ-ABC Distance Analysis", fontsize=14, fontweight='bold')
+        
+        # Distance distribution histogram
+        ax6 = fig.add_subplot(gs[1, 2:])  # Span 2 columns
+        
+        tooth_distances = self.create_distance_heatmap(results)
+        if tooth_distances:
+            all_vals = []
+            labels = []
+            colors = plt.cm.Set3(np.linspace(0, 1, len(tooth_distances)))
+            
+            for i, (tooth, distances) in enumerate(tooth_distances.items()):
+                ax6.hist(distances, bins=20, alpha=0.6, label=tooth, color=colors[i])
+                all_vals.extend(distances)
+            
+            ax6.set_xlabel("Distance (pixels)", fontsize=12, fontweight='bold')
+            ax6.set_ylabel("Frequency", fontsize=12, fontweight='bold')
+            ax6.set_title("Distance Distribution Across All Measurements", fontsize=14, fontweight='bold')
+            ax6.legend()
+            ax6.grid(True, alpha=0.3)
+        else:
+            ax6.text(0.5, 0.5, "No distance data available for histogram", 
+                    ha='center', va='center', transform=ax6.transAxes, fontsize=14)
+            ax6.set_title("Distance Distribution", fontsize=14, fontweight='bold')
+        
+        # Bottom row - Summary statistics table
+        ax7 = fig.add_subplot(gs[2, :])
+        ax7.axis('tight')
+        ax7.axis('off')
+        
+        # Create summary table
+        if tooth_means:
+            table_data = []
+            headers = ['Tooth', 'Mean Distance (px)', 'Max Distance (px)', 'Min Distance (px)', 'Range (px)', 'Measurements']
+            
+            for tooth_data in results["distance_analyses"]:
+                if tooth_data['analysis'] is not None:
+                    analysis = tooth_data['analysis']
+                    range_val = analysis['max_distance'] - analysis['min_distance']
+                    num_measurements = len(analysis['distances'])
+                    
+                    table_data.append([
+                        f"T{tooth_data['tooth_id']}",
+                        f"{analysis['mean_distance']:.2f}",
+                        f"{analysis['max_distance']:.2f}",
+                        f"{analysis['min_distance']:.2f}",
+                        f"{range_val:.2f}",
+                        str(num_measurements)
+                    ])
+            
+            table = ax7.table(cellText=table_data, colLabels=headers, 
+                             cellLoc='center', loc='center')
+            table.auto_set_font_size(False)
+            table.set_fontsize(10)
+            table.scale(1, 2)
+            
+            # Style the table
+            for (i, j), cell in table.get_celld().items():
+                if i == 0:  # Header row
+                    cell.set_text_props(weight='bold', color='white')
+                    cell.set_facecolor('#4472C4')
+                else:
+                    cell.set_facecolor('#F2F2F2' if i % 2 == 0 else 'white')
+                    
+            ax7.set_title("Detailed Distance Analysis Summary", fontsize=14, fontweight='bold', pad=20)
+        else:
+            ax7.text(0.5, 0.5, "No measurements available for summary table", 
+                    ha='center', va='center', transform=ax7.transAxes, fontsize=14)
+        
+        plt.suptitle("Comprehensive Dental CEJ-ABC Distance Analysis", fontsize=16, fontweight='bold', y=0.98)
+        
+        if save_path:
+            plt.savefig(save_path, dpi=300, bbox_inches="tight", facecolor='white')
+            print(f"Saved enhanced visualization to {save_path}")
+        
+        return fig
+
+
+if __name__ == "__main__":
+    unet_model = "unet.keras"
+    yolo_model = "yolov8n-seg.pt"
+    image_path = "trial.jpg"
+
+    seg = SimpleDentalSegmentationNoEnhance(unet_model, yolo_model)
+    res = seg.analyze_image(image_path)
+    
+    # Print distance analysis summary
+    seg.print_distance_summary(res)
+    
+    fig = seg.visualize_results(res, save_path="segmentation_with_distances.png")
     plt.show()