Update periodontitis_detection.py

periodontitis_detection.py

@@ -5,85 +5,75 @@ 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
+    def __init__(self, unet_model_path, yolo_model_path, unet_input_size=(224, 224, 3)):
+        """
+        Initialize the dental segmentation and analysis pipeline.
+        """
+        # Load Keras U-Net model
+        self.unet = tf.keras.models.load_model(unet_model_path)
         self.in_h, self.in_w, self.in_c = unet_input_size
 
-        # Load YOLOv8
+        # Load YOLOv8 (PyTorch) model
         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']}")
+        print("Models loaded successfully.")
+        print(f"Keras U-Net input shape: {self.unet.input_shape}")
+        print(f"Keras U-Net output shape: {self.unet.output_shape}")
+        print(f"YOLO model loaded: {yolo_model_path}")
 
     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
+        """
+        Prepare a BGR image for U-Net prediction.
+        Converts to RGB, resizes, and normalizes.
+        """
+        img_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+        img_resized = cv2.resize(img_rgb, (self.in_w, self.in_h), interpolation=cv2.INTER_LINEAR)
+        img_norm = img_resized.astype(np.float32) / 255.0
+        input_tensor = np.expand_dims(img_norm, axis=0)
+        return input_tensor, img_resized
 
     def run_unet(self, image_bgr):
+        """
+        Run the Keras U-Net model on the given image.
+        Returns CEJ and ABC masks.
+        """
         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]
+        preds = self.unet.predict(input_tensor, verbose=0)
+        out = preds[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)
+        elif out.ndim == 3 and out.shape[2] == 1:
+            binary = out[:, :, 0]
+            abc = (binary > 0.5).astype(np.uint8)
+            cej = np.zeros_like(abc)
         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)
+            h, w = out.shape[:2]
+            abc = np.zeros((h, w), dtype=np.uint8)
+            cej = np.zeros((h, w), dtype=np.uint8)
 
         return cej, abc, model_resized_image
 
     def detect_teeth(self, image_bgr):
+        """
+        Detect teeth using YOLOv8 PyTorch model.
+        Returns bounding boxes and confidence scores.
+        """
         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)
+            for box in boxes:
+                xyxy = box.xyxy[0].cpu().numpy()
+                conf = float(box.conf[0].cpu().numpy())
                 detections.append({
                     "bbox": xyxy.astype(np.float32),
                     "confidence": conf,
@@ -92,555 +82,198 @@ class SimpleDentalSegmentationNoEnhance:
         return detections
 
     def resize_mask_to_original(self, mask, original_shape):
+        """Resize a predicted mask back to original image size."""
         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
+        mask_resized = cv2.resize(mask.astype(np.uint8) * 255, (w_orig, h_orig), interpolation=cv2.INTER_NEAREST)
+        return (mask_resized > 127).astype(np.uint8)
 
     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)
+        """Extract the uppermost ABC line within a detected tooth bounding box."""
+        x1, y1, x2, y2 = map(int, bbox)
         height, width = abc_mask.shape
-        x1 = max(0, x1)
-        y1 = max(0, y1)
-        x2 = min(width-1, x2)
-        y2 = min(height-1, y2)
-        
+        x1, y1 = max(0, x1), max(0, y1)
+        x2, y2 = min(width - 1, x2), 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
+        for x in range(x1, x2 + 1):
+            column = np.where(abc_mask[y1:y2 + 1, x] == 1)[0]
+            if len(column) > 0:
+                y_absolute = y1 + np.min(column)
                 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)
+        """Extract the lowermost CEJ line within a detected tooth bounding box."""
+        x1, y1, x2, y2 = map(int, bbox)
         height, width = cej_mask.shape
-        x1 = max(0, x1)
-        y1 = max(0, y1)
-        x2 = min(width-1, x2)
-        y2 = min(height-1, y2)
-        
+        x1, y1 = max(0, x1), max(0, y1)
+        x2, y2 = min(width - 1, x2), 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
+        for x in range(x1, x2 + 1):
+            column = np.where(cej_mask[y1:y2 + 1, x] == 1)[0]
+            if len(column) > 0:
+                y_absolute = y1 + np.max(column)
                 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):
+        """Smooth a polyline using a simple moving average."""
         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)
+        pts = points.reshape(-1, 2)
+        smoothed = []
+        for i in range(len(pts)):
+            start, end = max(0, i - window_size // 2), min(len(pts), i + window_size // 2 + 1)
+            smoothed_y = np.mean(pts[start:end, 1])
+            smoothed.append([pts[i, 0], smoothed_y])
+        return np.array(smoothed, 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.
-        """
+        """Compute vertical distances between CEJ and ABC points."""
         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
+
+        cej_2d, abc_2d = cej_points.reshape(-1, 2), abc_points.reshape(-1, 2)
+        cej_dict = {x: y for x, y in cej_2d}
+        abc_dict = {x: y for x, y in abc_2d}
+
         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 = []
-        
+
+        distances, connections = [], []
         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)])
-        
+            cej_y, abc_y = cej_dict[x], abc_dict[x]
+            dist = abs(abc_y - cej_y)
+            distances.append({'x': x, 'cej_y': cej_y, 'abc_y': abc_y, 'distance': dist})
+            connections.append([(x, cej_y), (x, abc_y)])
+
         return {
             'distances': distances,
-            'connection_points': connection_points,
+            'connection_points': connections,
             '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])
+            'max_distance': np.max([d['distance'] for d in distances]),
+            'min_distance': np.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
+        """Interpolate CEJ and ABC lines when x-coordinates don’t match exactly."""
         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)
-        
+        x_min, x_max = max(cej_x_min, abc_x_min), 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))])
-        
+
+        x_samples = np.linspace(x_min, x_max, min(50, int(x_max - x_min) + 1), dtype=int)
+        cej_y = np.interp(x_samples, cej_points[:, 0], cej_points[:, 1])
+        abc_y = np.interp(x_samples, abc_points[:, 0], abc_points[:, 1])
+
+        distances, connections = [], []
+        for x, cy, ay in zip(x_samples, cej_y, abc_y):
+            dist = abs(ay - cy)
+            distances.append({'x': int(x), 'cej_y': int(cy), 'abc_y': int(ay), 'distance': dist})
+            connections.append([(int(x), int(cy)), (int(x), int(ay))])
+
         return {
             'distances': distances,
-            'connection_points': connection_points,
+            'connection_points': connections,
             '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])
+            'max_distance': np.max([d['distance'] for d in distances]),
+            'min_distance': np.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.
-        """
+        """Draw color-coded CEJ-ABC measurement lines."""
         if distance_analysis is None:
             return image
-            
-        img_with_distances = image.copy()
-        
-        # Draw connection lines with gradient effect
-        connection_points = distance_analysis['connection_points']
+        img = image.copy()
+        connections = 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
+            return img
+
+        min_d, max_d = min(distances), max(distances)
+        dist_range = max_d - min_d if max_d != min_d else 1
+
+        for i in range(0, len(connections), 3):  # draw every 3rd to reduce clutter
+            (x1, y1), (x2, y2) = connections[i]
+            dist = distances[i]
+            norm = (dist - min_d) / dist_range
+            color = (0, int(255 * (1 - norm)), int(255 * norm))  # Green→Red
+            cv2.line(img, (x1, y1), (x2, y2), color, max(1, int(2 + 2 * norm)))
+        return img
 
     def analyze_image(self, image_path):
+        """
+        Perform full analysis on a dental image:
+        segmentation, detection, distance measurement, and visualization.
+        """
         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]
 
+        h_orig, w_orig = img_bgr.shape[:2]
         cej_unet, abc_unet, _ = self.run_unet(img_bgr)
+        cej_mask = self.resize_mask_to_original(cej_unet, (h_orig, w_orig))
+        abc_mask = self.resize_mask_to_original(abc_unet, (h_orig, w_orig))
+        detections = self.detect_teeth(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))
+        combined = img_bgr.copy()
+        all_results = []
 
-        cej_bin = (cej_orig > 0.5).astype(np.uint8) * 255
-        abc_bin = (abc_orig > 0.5).astype(np.uint8) * 255
+        for det in detections:
+            x1, y1, x2, y2 = det["bbox"].astype(int)
+            cv2.rectangle(combined, (x1, y1), (x2, y2), (0, 255, 0), 2)
 
-        detections = self.detect_teeth(img_bgr)
+            abc_line = self.extract_abc_uppermost_line_within_bbox(abc_mask, (x1, y1, x2, y2))
+            cej_line = self.extract_cej_lowermost_line_within_bbox(cej_mask, (x1, y1, x2, y2))
 
-        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
+            if abc_line is not None:
+                abc_line = self.smooth_landmarks(abc_line)
+                cv2.polylines(combined, [abc_line], False, (255, 0, 0), 2)
+            if cej_line is not None:
+                cej_line = self.smooth_landmarks(cej_line)
+                cv2.polylines(combined, [cej_line], False, (0, 0, 255), 2)
 
-            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
+            if cej_line is not None and abc_line is not None:
+                distance_analysis = self.compute_cej_abc_distances(cej_line, abc_line)
+                if distance_analysis:
+                    combined = self.draw_distance_measurements(combined, distance_analysis, det["tooth_id"])
+
+            all_results.append({
+                "tooth_id": det["tooth_id"],
+                "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 = {
+        return {
             "original": img_bgr,
-            "cej_mask": cej_bin,
-            "abc_mask": abc_bin,
+            "cej_mask": cej_mask,
+            "abc_mask": abc_mask,
             "detections": detections,
             "combined": combined,
-            "abc_segments": all_abc_segments,
-            "cej_segments": all_cej_segments,
-            "distance_analyses": all_distance_analyses
+            "distance_analyses": all_results
         }
-        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"
+    unet_model = "unet.keras"  # Keras model
+    yolo_model = "yolov8n-seg.pt"               # YOLOv8 PyTorch model
+    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")
-    plt.show()
+
+    plt.figure(figsize=(12, 8))
+    plt.imshow(cv2.cvtColor(res["combined"], cv2.COLOR_BGR2RGB))
+    plt.title("Dental CEJ–ABC Analysis Result")
+    plt.axis("off")
+    plt.show()