Dental Width Analyser

dental-width-analyzer.py

#!/usr/bin/env python3
"""
DenteScope AI - Tooth Width Measurement System
Measures width differences between Primary Second Molars and Second Premolars
For orthodontic space analysis on panoramic dental X-rays
"""

import cv2
import numpy as np
from ultralytics import YOLO
from pathlib import Path
import json

class ToothWidthAnalyzer:
    """
    AI-powered tooth width measurement for orthodontic analysis
    """
    
    # FDI Notation for teeth we care about
    TEETH_OF_INTEREST = {
        # Second Primary Molars (Deciduous)
        'primary_molars': {
            55: 'Q1-Primary-Molar-2',  # Upper right
            65: 'Q2-Primary-Molar-2',  # Upper left
            75: 'Q3-Primary-Molar-2',  # Lower left
            85: 'Q4-Primary-Molar-2',  # Lower right
        },
        # Second Premolars (Permanent)
        'premolars': {
            15: 'Q1-Premolar-2',  # Upper right
            25: 'Q2-Premolar-2',  # Upper left
            35: 'Q3-Premolar-2',  # Lower left
            45: 'Q4-Premolar-2',  # Lower right
        }
    }
    
    def __init__(self, model_path='yolov8n.pt', calibration_ref_mm=None):
        """
        Initialize tooth width analyzer
        
        Args:
            model_path: Path to trained YOLOv8 model
            calibration_ref_mm: Known reference size in mm for calibration
        """
        self.model = YOLO(model_path)
        self.calibration_ref_mm = calibration_ref_mm
        self.pixel_to_mm_ratio = None
    
    def detect_calibration_marker(self, image):
        """
        Detect calibration marker or use known tooth size for pixel-to-mm conversion
        For panoramic X-rays, we can use average tooth width as reference
        
        Average adult second premolar width: ~7mm
        """
        # TODO: Implement automatic calibration detection
        # For now, use estimated ratio based on image resolution
        height, width = image.shape[:2]
        
        # Typical panoramic X-ray shows ~120mm width of dental arch
        # Assuming image width represents approximately 120mm
        estimated_mm_width = 120
        self.pixel_to_mm_ratio = estimated_mm_width / width
        
        return self.pixel_to_mm_ratio
    
    def detect_teeth(self, image):
        """
        Detect all teeth in the image using YOLOv8
        
        Returns:
            List of detected teeth with bounding boxes and classifications
        """
        results = self.model(image)
        detections = []
        
        for result in results:
            boxes = result.boxes
            for box in boxes:
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                confidence = float(box.conf[0])
                class_id = int(box.cls[0])
                
                detection = {
                    'bbox': (x1, y1, x2, y2),
                    'confidence': confidence,
                    'class_id': class_id,
                    'type': self._classify_tooth_type(class_id),
                }
                
                detections.append(detection)
        
        return detections
    
    def _classify_tooth_type(self, class_id):
        """
        Classify tooth as primary molar or premolar
        """
        # Map class_id to tooth type
        # This depends on your training data labels
        if class_id in [4, 5, 20, 21, 28, 29]:  # Primary molars (positions 5,6 in each quadrant)
            return 'primary_molar'
        elif class_id in [4, 12, 20, 28]:  # Second premolars (position 5 in each quadrant)
            return 'premolar'
        return 'other'
    
    def measure_tooth_width(self, image, bbox, calibration_ratio=None):
        """
        Measure tooth width from bounding box
        
        Args:
            image: Input image
            bbox: Bounding box (x1, y1, x2, y2)
            calibration_ratio: Pixel to mm conversion ratio
            
        Returns:
            Width in mm
        """
        x1, y1, x2, y2 = bbox
        
        # Calculate width in pixels
        width_pixels = x2 - x1
        
        # Convert to mm
        if calibration_ratio:
            width_mm = width_pixels * calibration_ratio
        elif self.pixel_to_mm_ratio:
            width_mm = width_pixels * self.pixel_to_mm_ratio
        else:
            # Use default calibration
            self.detect_calibration_marker(image)
            width_mm = width_pixels * self.pixel_to_mm_ratio
        
        return round(width_mm, 2)
    
    def analyze_width_discrepancy(self, primary_width, premolar_width):
        """
        Calculate width discrepancy and provide clinical interpretation
        
        Args:
            primary_width: Width of primary molar in mm
            premolar_width: Width of premolar in mm
            
        Returns:
            Analysis dictionary with discrepancy and recommendations
        """
        discrepancy = primary_width - premolar_width
        
        analysis = {
            'primary_molar_width_mm': primary_width,
            'premolar_width_mm': premolar_width,
            'discrepancy_mm': round(discrepancy, 2),
            'clinical_significance': self._interpret_discrepancy(discrepancy)
        }
        
        return analysis
    
    def _interpret_discrepancy(self, discrepancy_mm):
        """
        Provide clinical interpretation of width discrepancy
        """
        if discrepancy_mm > 2:
            return {
                'level': 'SIGNIFICANT',
                'color': (0, 0, 255),  # Red
                'recommendation': 'Space maintainer recommended. Primary molar significantly wider than successor.',
                'space_loss_risk': 'HIGH'
            }
        elif discrepancy_mm > 0.5:
            return {
                'level': 'MODERATE',
                'color': (0, 165, 255),  # Orange
                'recommendation': 'Monitor space. Consider space maintainer if early loss.',
                'space_loss_risk': 'MODERATE'
            }
        elif discrepancy_mm > -0.5:
            return {
                'level': 'MINIMAL',
                'color': (0, 255, 255),  # Yellow
                'recommendation': 'Normal variation. Standard monitoring.',
                'space_loss_risk': 'LOW'
            }
        else:
            return {
                'level': 'FAVORABLE',
                'color': (0, 255, 0),  # Green
                'recommendation': 'Successor tooth larger. Adequate space expected.',
                'space_loss_risk': 'VERY LOW'
            }
    
    def visualize_measurements(self, image, detections, measurements):
        """
        Visualize tooth detections and width measurements on image
        """
        viz_image = image.copy()
        
        for detection, measurement in zip(detections, measurements):
            bbox = detection['bbox']
            x1, y1, x2, y2 = bbox
            tooth_type = detection['type']
            
            # Color code by tooth type
            if tooth_type == 'primary_molar':
                color = (255, 0, 255)  # Magenta for primary
                label = f"Primary Molar: {measurement['width_mm']:.2f}mm"
            elif tooth_type == 'premolar':
                color = (0, 255, 255)  # Cyan for premolar
                label = f"Premolar: {measurement['width_mm']:.2f}mm"
            else:
                color = (128, 128, 128)  # Gray for others
                label = f"Tooth: {measurement['width_mm']:.2f}mm"
            
            # Draw bounding box
            cv2.rectangle(viz_image, (x1, y1), (x2, y2), color, 2)
            
            # Draw width measurement line
            mid_y = (y1 + y2) // 2
            cv2.line(viz_image, (x1, mid_y), (x2, mid_y), color, 2)
            cv2.circle(viz_image, (x1, mid_y), 3, color, -1)
            cv2.circle(viz_image, (x2, mid_y), 3, color, -1)
            
            # Add label
            cv2.putText(viz_image, label, (x1, y1-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
        
        return viz_image
    
    def generate_report(self, image_path, detections, analyses):
        """
        Generate comprehensive orthodontic report
        """
        report = {
            'image': str(image_path),
            'analysis_date': str(Path.cwd()),
            'findings': [],
            'summary': {
                'total_primary_molars': 0,
                'total_premolars': 0,
                'high_risk_quadrants': [],
                'recommendations': []
            }
        }
        
        # Analyze each quadrant
        for quadrant in ['Q1', 'Q2', 'Q3', 'Q4']:
            primary = None
            premolar = None
            
            # Find teeth in this quadrant
            for detection, analysis in zip(detections, analyses):
                if quadrant in detection.get('fdi_label', ''):
                    if detection['type'] == 'primary_molar':
                        primary = analysis
                    elif detection['type'] == 'premolar':
                        premolar = analysis
            
            # Generate quadrant analysis
            if primary and premolar:
                discrepancy_analysis = self.analyze_width_discrepancy(
                    primary['width_mm'], 
                    premolar['width_mm']
                )
                
                report['findings'].append({
                    'quadrant': quadrant,
                    **discrepancy_analysis
                })
                
                if discrepancy_analysis['clinical_significance']['space_loss_risk'] in ['HIGH', 'MODERATE']:
                    report['summary']['high_risk_quadrants'].append(quadrant)
        
        # Generate recommendations
        if report['summary']['high_risk_quadrants']:
            report['summary']['recommendations'].append(
                f"Space maintainers recommended for quadrants: {', '.join(report['summary']['high_risk_quadrants'])}"
            )
        
        return report


def main():
    """
    Example usage for event demo
    """
    # Initialize analyzer
    analyzer = ToothWidthAnalyzer(
        model_path='runs/train/tooth_detection/weights/best.pt'
    )
    
    # Load test image
    test_image_path = "data/raw/test_panoramic.jpg"
    image = cv2.imread(test_image_path)
    
    if image is None:
        print(f"❌ Could not load image: {test_image_path}")
        return
    
    print("🔍 Detecting teeth...")
    detections = analyzer.detect_teeth(image)
    
    print("📏 Measuring tooth widths...")
    measurements = []
    for detection in detections:
        width_mm = analyzer.measure_tooth_width(image, detection['bbox'])
        measurements.append({
            'width_mm': width_mm,
            'type': detection['type']
        })
    
    print("📊 Generating analysis...")
    # Visualize results
    viz_image = analyzer.visualize_measurements(image, detections, measurements)
    
    # Save result
    output_path = "tooth_width_analysis_result.jpg"
    cv2.imwrite(output_path, viz_image)
    
    print(f"\n✅ Analysis complete!")
    print(f"📸 Result saved to: {output_path}")
    
    # Generate report
    report = analyzer.generate_report(test_image_path, detections, measurements)
    
    # Save report as JSON
    with open('tooth_width_report.json', 'w') as f:
        json.dump(report, f, indent=2)
    
    print(f"📄 Report saved to: tooth_width_report.json")
    
    # Display summary
    print("\n" + "="*60)
    print("📋 TOOTH WIDTH ANALYSIS SUMMARY")
    print("="*60)
    for finding in report.get('findings', []):
        print(f"\nQuadrant {finding['quadrant']}:")
        print(f"  Primary Molar: {finding['primary_molar_width_mm']:.2f}mm")
        print(f"  Premolar: {finding['premolar_width_mm']:.2f}mm")
        print(f"  Discrepancy: {finding['discrepancy_mm']:.2f}mm")
        print(f"  Risk Level: {finding['clinical_significance']['level']}")
        print(f"  Recommendation: {finding['clinical_significance']['recommendation']}")


if __name__ == "__main__":
    main()