omr_detector.py

import cv2
import numpy as np
from typing import List, Tuple, Dict
import json
import os

class OMRDetector:
    def __init__(self, image_path: str, crop_bottom_percent: float = 60, crop_top_range: Tuple[float, float] = None):
        """
        Initialize OMR Detector with image path
        
        Args:
            image_path: Path to the OMR sheet image
            crop_bottom_percent: Percentage of image to crop from bottom (default 60%)
            crop_top_range: Tuple (start_percent, end_percent) to crop top section (e.g., (20, 40) for header info)
        """
        self.image_path = image_path
        self.original = cv2.imread(image_path)
        if self.original is None:
            raise ValueError(f"Could not read image from {image_path}")
        
        # Store full image dimensions
        self.full_height, self.full_width = self.original.shape[:2]
        
        # Calculate crop region (bottom X% of image or top section)
        if crop_top_range:
            # Crop top section (for header information like roll, class, etc.)
            start_percent, end_percent = crop_top_range
            crop_start_y = int(self.full_height * start_percent / 100)
            crop_end_y = int(self.full_height * end_percent / 100)
            self.cropped_original = self.original[crop_start_y:crop_end_y, :]
            self.crop_offset_y = crop_start_y
        else:
            # Crop bottom section (for answers)
            self.crop_percent = crop_bottom_percent
            crop_start_y = int(self.full_height * (1 - crop_bottom_percent / 100))
            self.cropped_original = self.original[crop_start_y:, :]
            self.crop_offset_y = crop_start_y
        
        # Convert cropped image to grayscale
        self.gray = cv2.cvtColor(self.cropped_original, cv2.COLOR_BGR2GRAY)
        self.height, self.width = self.gray.shape
        
    def preprocess_image(self, save_debug=False) -> np.ndarray:
        """
        Preprocess the image for better circle detection
        
        Args:
            save_debug: If True, save debug binary image
        
        Returns:
            Processed binary image
        """
        # Apply Gaussian blur to reduce noise
        blurred = cv2.GaussianBlur(self.gray, (3, 3), 0)
        
        # Try adaptive thresholding first for better local contrast
        binary = cv2.adaptiveThreshold(
            blurred, 255,
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY_INV,
            15, 3
        )
        
        # Apply light morphological closing to fill small gaps
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2))
        binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel, iterations=1)
        
        # Save debug image
        if save_debug:
            debug_path = self.image_path.replace(".jpg", "_binary_debug.jpg")
            cv2.imwrite(debug_path, binary)
        
        return binary
    
    def detect_circles(self, binary_image: np.ndarray) -> List[Tuple[int, int, int]]:
        """
        Detect circles in the binary image using contour analysis
        
        Args:
            binary_image: Preprocessed binary image
            
        Returns:
            List of circles as (x, y, radius)
        """
        # Find contours - use RETR_LIST to get all contours including inner ones
        contours, _ = cv2.findContours(
            binary_image, 
            cv2.RETR_LIST, 
            cv2.CHAIN_APPROX_SIMPLE
        )
        
        circles = []
        
        for contour in contours:
            # Calculate contour properties
            area = cv2.contourArea(contour)
            
            # More relaxed area filter to catch all circles
            if area < 250 or area > 4500:
                continue
            
            perimeter = cv2.arcLength(contour, True)
            
            if perimeter == 0:
                continue
            
            # Calculate circularity (0-1, where 1 is perfect circle)
            circularity = 4 * np.pi * area / (perimeter * perimeter)
            
            # More relaxed circularity to catch slightly distorted circles
            if circularity < 0.50:
                continue
            
            # Get bounding circle
            (x, y), radius = cv2.minEnclosingCircle(contour)
            
            # Filter by radius (more relaxed range)
            if 18 <= radius <= 36:
                circles.append((int(x), int(y), int(radius)))
        
        # Remove duplicate circles (same position, different radius due to inner/outer contours)
        unique_circles = []
        for circle in circles:
            x, y, r = circle
            # Check if this circle is too close to an existing one
            is_duplicate = False
            for existing in unique_circles:
                ex, ey, er = existing
                distance = np.sqrt((x - ex)**2 + (y - ey)**2)
                if distance < 15:  # If centers are within 15 pixels, consider duplicate
                    # Keep the one with larger radius (outer contour)
                    if r > er:
                        unique_circles.remove(existing)
                        unique_circles.append(circle)
                    is_duplicate = True
                    break
            
            if not is_duplicate:
                unique_circles.append(circle)
        
        return unique_circles
    
    def is_circle_filled(self, circle: Tuple[int, int, int], threshold: float = 0.5) -> bool:
        """
        Check if a circle is filled/marked with improved accuracy
        
        Args:
            circle: Circle as (x, y, radius)
            threshold: Fill ratio threshold (0-1)
            
        Returns:
            True if circle is filled, False otherwise
        """
        x, y, radius = circle
        
        # Create a mask for the circle (use smaller radius to avoid borders)
        mask = np.zeros(self.gray.shape, dtype=np.uint8)
        inner_radius = max(int(radius * 0.7), radius - 5)  # Check inner 70% of circle
        cv2.circle(mask, (x, y), inner_radius, 255, -1)
        
        # Get the pixels within the circle
        circle_pixels = cv2.bitwise_and(self.gray, self.gray, mask=mask)
        
        # Calculate the mean intensity
        masked_pixels = circle_pixels[mask == 255]
        
        if len(masked_pixels) == 0:
            return False
        
        mean_intensity = np.mean(masked_pixels)
        std_intensity = np.std(masked_pixels)
        
        # Calculate the percentage of dark pixels (below threshold)
        dark_pixel_threshold = 120
        dark_pixels = np.sum(masked_pixels < dark_pixel_threshold)
        total_pixels = len(masked_pixels)
        dark_ratio = dark_pixels / total_pixels if total_pixels > 0 else 0
        
        # A circle is filled if:
        # 1. Mean intensity is low (dark) AND
        # 2. High percentage of dark pixels (>50%) AND
        # 3. Low standard deviation (uniformly filled)
        is_filled = (mean_intensity < 140) and (dark_ratio > 0.5) and (std_intensity < 60)
        
        return is_filled
    
    def group_circles_into_grid(self, circles: List[Tuple[int, int, int]], options_per_question: int = 4) -> Dict:
        """
        Group circles into columns and organize by questions (4 options per question)
        Question numbering: top to bottom in each column, then left to right
        
        Args:
            circles: List of detected circles
            options_per_question: Number of options per question (default: 4)
            
        Returns:
            Dictionary with organized grid structure
        """
        if not circles:
            return {}
        
        # Sort circles by y-coordinate (top to bottom), then x-coordinate (left to right)
        sorted_circles = sorted(circles, key=lambda c: (c[1], c[0]))
        
        # Group circles by rows (similar y-coordinates)
        rows = []
        current_row = [sorted_circles[0]]
        row_tolerance = 35  # Pixels tolerance for same row
        
        for circle in sorted_circles[1:]:
            if abs(circle[1] - current_row[0][1]) < row_tolerance:
                current_row.append(circle)
            else:
                rows.append(sorted(current_row, key=lambda c: c[0]))
                current_row = [circle]
        
        if current_row:
            rows.append(sorted(current_row, key=lambda c: c[0]))
        
        # Convert rows to column-major order for Y-wise numbering
        # First, organize into a 2D grid
        max_cols = max(len(row) for row in rows) if rows else 0
        
        # Group by columns (vertical sections of 4 options each)
        questions = []
        incomplete_questions = []
        col_groups = max_cols // options_per_question  # Number of question columns
        
        for col_group in range(col_groups):
            # For each column group (represents multiple questions vertically)
            start_col = col_group * options_per_question
            end_col = start_col + options_per_question
            
            # Extract this column group from all rows
            for row_idx, row in enumerate(rows):
                if len(row) > start_col:
                    question_options = row[start_col:end_col]
                    if len(question_options) == options_per_question:
                        questions.append(question_options)
                    elif len(question_options) > 0:
                        # Incomplete question - pad with None to maintain position
                        incomplete_questions.append({
                            'position': len(questions),
                            'found': len(question_options),
                            'expected': options_per_question,
                            'row': row_idx + 1,
                            'col_group': col_group + 1
                        })
                        # Add as incomplete but still count it
                        # Pad with None placeholders
                        padded_options = question_options + [None] * (options_per_question - len(question_options))
                        questions.append(padded_options)
        
        # Report incomplete questions (silently handled)
        
        return questions
    
    def detect_answers(self, radius_filter: tuple = None) -> List[Dict]:
        """
        Main function to detect marked answers in OMR sheet
        
        Args:
            radius_filter: Tuple (min_radius, max_radius) to filter circles by radius range
        
        Returns:
            List of dictionaries with question numbers and marked options
        """
        # Preprocess image
        binary = self.preprocess_image(save_debug=False)
        
        # Detect circles
        all_circles = self.detect_circles(binary)
        
        # Apply radius filter if specified
        if radius_filter:
            min_r, max_r = radius_filter
            circles = [c for c in all_circles if min_r <= c[2] <= max_r]
        else:
            circles = all_circles
        
        # Group circles into questions (4 options each)
        questions = self.group_circles_into_grid(circles, options_per_question=4)
        
        # Analyze each question to find marked option
        results = {}
        
        for question_num, question_options in enumerate(questions, 1):
            # Skip if question has None placeholders (incomplete)
            if None in question_options:
                continue
            
            # Check which option is marked (if any)
            for option_idx, circle in enumerate(question_options):
                x, y, radius = circle
                if self.is_circle_filled(circle):
                    results[str(question_num)] = option_idx + 1  # Option 1, 2, 3, or 4
                    break  # Only one option should be marked per question
        
        return results
    
    def detect_header_info(self) -> Dict:
        """
        Detect header information (serial, roll, class, subject code, set code)
        from the top section of the OMR sheet (20-40% height)
        
        Returns:
            Dictionary containing serial, roll, class, subject_code, and set_code
        """
        # Preprocess image
        binary = self.preprocess_image(save_debug=False)
        
        # Detect all circles in the header section
        all_circles = self.detect_circles(binary)
        
        # Filter circles by radius (header circles might be slightly different size)
        circles = [c for c in all_circles if 18 <= c[2] <= 36]
        
        # Sort circles by position for easier processing
        sorted_circles = sorted(circles, key=lambda c: (c[0], c[1]))  # Sort by x, then y
        
        # Detect filled circles and their positions
        filled_circles = []
        for circle in sorted_circles:
            if self.is_circle_filled(circle):
                filled_circles.append(circle)
        
        # Group circles into columns (for different fields)
        header_data = self._extract_header_fields(sorted_circles, filled_circles)
        
        return header_data
    
    def _extract_header_fields(self, all_circles: List[Tuple[int, int, int]], 
                                filled_circles: List[Tuple[int, int, int]]) -> Dict:
        """
        Extract header fields from detected circles
        
        Args:
            all_circles: All detected circles sorted by position
            filled_circles: Only filled/marked circles
            
        Returns:
            Dictionary with serial, roll, class, subject_code, set_code
        """
        # Group circles by X position (columns)
        columns = []
        current_column = []
        x_tolerance = 50  # Pixels tolerance for same column
        
        for circle in all_circles:
            if not current_column:
                current_column.append(circle)
            else:
                # Check if this circle is in the same column
                avg_x = sum(c[0] for c in current_column) / len(current_column)
                if abs(circle[0] - avg_x) < x_tolerance:
                    current_column.append(circle)
                else:
                    # Sort current column by Y and add to columns
                    columns.append(sorted(current_column, key=lambda c: c[1]))
                    current_column = [circle]
        
        if current_column:
            columns.append(sorted(current_column, key=lambda c: c[1]))
        
        # Initialize result
        result = {
            'serial': None,
            'roll': None,
            'class': None,
            'subject_code': None,
            'set_code': None
        }
        
        # Analyze each column to determine which field it represents
        # Based on the actual image structure:
        # Column 0: Serial/Class (7-8 circles)
        # Columns 1-6: Roll number (6 digits, 10 circles each = digits 0-9)
        # Columns 7-9: Subject code (3 digits, 10 circles each)
        # Column 10: Set code (4-6 circles for Bengali letters ক খ গ ঘ)
        
        # Track which columns are processed
        roll_digits = []
        subject_digits = []
        
        for col_idx, column in enumerate(columns):
            # Find which circle is filled in this column
            filled_in_column = [c for c in filled_circles if any(
                abs(c[0] - circle[0]) < 20 and abs(c[1] - circle[1]) < 20 
                for circle in column
            )]
            
            if filled_in_column:
                # Find position of filled circle in this column
                filled_circle = filled_in_column[0]
                position = 0
                for i, circle in enumerate(column):
                    if abs(filled_circle[0] - circle[0]) < 20 and abs(filled_circle[1] - circle[1]) < 20:
                        position = i
                        break
                
                # Determine field based on column structure
                if col_idx == 0:
                    # First column - could be serial or class
                    if len(column) <= 8:
                        # Smaller column = Serial/Class
                        # The serial column often has values like 10, 20, 30, etc.
                        # or 01-99. For 7 circles, it might represent: ওয়েন, ৩, ৬, etc.
                        # Position mapping may vary, but typically direct
                        result['serial'] = position + 1  # 1-indexed for serial numbers
                    else:
                        # 10 circles = first digit of roll
                        roll_digits.append(self._map_position_to_digit(position, column))
                elif len(column) == 10 and len(roll_digits) < 6:
                    # 10 circles and roll not complete = Roll digit
                    roll_digits.append(self._map_position_to_digit(position, column))
                elif len(column) == 10 and len(roll_digits) >= 6:
                    # 10 circles after roll complete = Subject code digit
                    subject_digits.append(self._map_position_to_digit(position, column))
                elif len(column) >= 4 and len(column) <= 6:
                    # 4-6 circles = Set code (Bengali letters)
                    # Position 0=ক, 1=খ, 2=গ, 3=ঘ, etc.
                    bengali_letters = ['ক', 'খ', 'গ', 'ঘ', 'ঙ', 'চ']
                    if position < len(bengali_letters):
                        result['set_code'] = bengali_letters[position]
                elif len(column) >= 2 and len(column) < 4:
                    # Very small column, possibly class indicator
                    # If we haven't set class yet, set it
                    if result['class'] is None:
                        # These small columns might indicate class level
                        # For now, just note that they exist
                        result['class'] = 'X'  # Placeholder - would need OCR or specific detection
        
        # Assemble roll number
        if roll_digits:
            result['roll'] = ''.join(str(d) for d in roll_digits)
        
        # Assemble subject code
        if subject_digits:
            result['subject_code'] = ''.join(str(d) for d in subject_digits)
        
        # Calculate class from serial number
        # Serial 1 = Class 6, Serial 2 = Class 7, ..., Serial 7 = Class 12
        if result['serial'] is not None:
            serial_num = result['serial']
            if 1 <= serial_num <= 7:
                result['class'] = str(5 + serial_num)  # Class 6-12
            elif serial_num == 0:
                result['class'] = '5'  # Serial 0 = Class 5 (if applicable)
        
        # Remove serial from response (only keep class)
        result.pop('serial', None)
        
        return result
    
    def _map_position_to_digit(self, position: int, column: List[Tuple[int, int, int]]) -> int:
        """
        Map circle position in column to digit value
        
        For OMR sheets, typically:
        - If 10 circles per column: positions 0-9 map to digits 0-9
        - Need to check the actual layout
        
        Args:
            position: Position of marked circle in column (0-based)
            column: All circles in the column
            
        Returns:
            Digit value (0-9)
        """
        # Based on the image, it appears the layout shows digits in sequence
        # You may need to adjust this mapping based on actual OMR layout
        
        # Common layouts:
        # Layout 1: Top to bottom = 0,1,2,3,4,5,6,7,8,9
        # Layout 2: Top to bottom = 1,2,3,4,5,6,7,8,9,0
        
        # From the image, it looks like the digits are arranged as: 2,4,6,8,0,2 (visible marks)
        # This suggests the layout might have pairs or specific arrangement
        
        # For now, assume direct mapping (may need adjustment)
        digit_map = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
        
        if position < len(digit_map):
            return digit_map[position]
        return 0
    
    def visualize_detection(self, output_path: str = "detected_omr.jpg", show_crop_region: bool = True, show_grid: bool = False):
        """
        Create a visualization of detected circles with color coding and grid
        
        Args:
            output_path: Path to save the output image
            show_crop_region: Whether to show the crop region boundary
            show_grid: Whether to show question grid boundaries and numbers
        """
        # Preprocess
        binary = self.preprocess_image()
        circles = self.detect_circles(binary)
        
        # Create output image from FULL original (not cropped)
        output = self.original.copy()
        
        # Get questions grid for visualization
        if show_grid:
            questions = self.group_circles_into_grid(circles, options_per_question=4)
            
            # Draw bounding boxes around each question and label them
            for question_num, question_options in enumerate(questions, 1):
                # Skip None placeholders
                valid_options = [opt for opt in question_options if opt is not None]
                if not valid_options:
                    continue
                
                # Get bounding box of all options in this question
                xs = [opt[0] for opt in valid_options]
                ys = [opt[1] for opt in valid_options]
                
                min_x, max_x = min(xs) - 40, max(xs) + 40
                min_y, max_y = min(ys) - 15, max(ys) + 15
                
                # Adjust for crop offset
                min_y_full = min_y + self.crop_offset_y
                max_y_full = max_y + self.crop_offset_y
                
                # Draw bounding box
                color = (255, 165, 0) if None in question_options else (0, 165, 255)  # Orange if incomplete, Blue if complete
                cv2.rectangle(output, (min_x, min_y_full), (max_x, max_y_full), color, 2)
                
                # Draw question number
                label_y = min_y_full - 5 if min_y_full > 30 else max_y_full + 20
                cv2.putText(output, f"Q{question_num}", 
                           (min_x, label_y),
                           cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)
        
        # Draw all circles with better color coding
        # Adjust circle coordinates to full image coordinates
        for idx, circle in enumerate(circles, 1):
            x, y, radius = circle
            # Adjust y coordinate for crop offset
            y_full = y + self.crop_offset_y
            
            if self.is_circle_filled(circle):
                # Draw filled circles with BRIGHT GREEN border only
                cv2.circle(output, (x, y_full), radius, (0, 255, 0), 4)
            else:
                # Draw empty circles with RED border only
                cv2.circle(output, (x, y_full), radius, (0, 0, 255), 3)
        
        # Save output
        cv2.imwrite(output_path, output)