time_table_extract.py

from pathlib import Path
from typing import Any, Dict, List, Tuple
import cv2
import numpy as np
import pytesseract
from pytesseract import Output


def _compute_bounds(
   proj: np.ndarray,
   threshold: float
) -> List[int]:
   """
   Extracts midpoint boundaries from a 1D projection array where values exceed threshold.
   """
   lines = np.nonzero(proj > threshold)[0]
   if lines.size == 0:
      return []

   bounds: List[int] = []
   start = int(lines[0])
   prev = start
   for idx in lines[1:]:
      if idx > prev + 1:
         bounds.append((start + prev) // 2)
         start = int(idx)
      prev = idx
   bounds.append((start + int(lines[-1])) // 2)
   return bounds


def extract_grid_cells(
   img: Any
) -> Tuple[
   Dict[int, Dict[int, Tuple[int, int, int, int]]],
   Dict[Tuple[int, int, int, int], Tuple[int, int]]
]:
   """
   Detects grid cells and returns:
      - row_cells: mapping row -> {col: rect}
      - cell_lookup: mapping rect -> (row, col)
   """
   ker_h = cv2.getStructuringElement(cv2.MORPH_RECT, (40, 1))
   ker_v = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 40))

   gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
   thresh = cv2.adaptiveThreshold(
      gray, 255,
      cv2.ADAPTIVE_THRESH_MEAN_C,
      cv2.THRESH_BINARY_INV,
      15, 10
   )

   horiz = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, ker_h)
   vert  = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, ker_v)

   v_proj = vert.sum(axis=0).astype(float)
   h_proj = horiz.sum(axis=1).astype(float)
   v_thresh = v_proj.max() * 0.5
   h_thresh = h_proj.max() * 0.5

   v_bounds = _compute_bounds(v_proj, v_thresh)
   h_bounds = _compute_bounds(h_proj, h_thresh)

   cell_lookup: Dict[Tuple[int, int, int, int], Tuple[int, int]] = {}
   row_cells: Dict[int, Dict[int, Tuple[int, int, int, int]]] = {}

   for r in range(len(h_bounds) - 1):
      row = r + 1
      y0, y1 = h_bounds[r], h_bounds[r + 1]
      row_cells[row] = {}
      for c in range(len(v_bounds) - 1):
         col = c + 1
         x0, x1 = v_bounds[c], v_bounds[c + 1]
         rect = (x0, y0, x1 - x0, y1 - y0)
         cell_lookup[rect] = (row, col)
         row_cells[row][col] = rect

   return row_cells, cell_lookup


def extract_day_venue(
   img: Any
) -> Tuple[Dict[int, str], Dict[int, str]]:
   """
   Extracts day (col 2) and venue (col 3) text based on OCR positions.
   """
   _, cell_lookup = extract_grid_cells(img)
   ocr = pytesseract.image_to_data(img, output_type=Output.DICT)
   positions: Dict[str, List[Tuple[int, int]]] = {}
   for i, raw in enumerate(ocr["text"]):
      text = raw.strip()
      if not text:
         continue
      x, y, w, h = (
         ocr["left"][i], ocr["top"][i],
         ocr["width"][i], ocr["height"][i]
      )
      cx, cy = x + w // 2, y + h // 2
      for rect, (r, c) in cell_lookup.items():
         bx, by, bw, bh = rect
         if bx <= cx <= bx + bw and by <= cy <= by + bh:
            positions.setdefault(text, []).append((r, c))
            break

   day_map: Dict[int, str] = {}
   venue_map: Dict[int, str] = {}
   for text, coords in positions.items():
      for r, c in coords:
         if c == 2:
            day_map[r] = text
         elif c == 3:
            venue_map[r] = text
         elif r == 4:
            room_map

   return day_map, venue_map


def extract_course_positions(
   img: Any,
   course_codes: List[str]
) -> Dict[str, List[Tuple[int, int]]]:
   """
   Finds row/col positions for given course codes via OCR.
   """
   _, cell_lookup = extract_grid_cells(img)
   ocr = pytesseract.image_to_data(img, output_type=Output.DICT)
   positions: Dict[str, List[Tuple[int, int]]] = {text.strip(): [] for text in ocr["text"]}

   for i, raw in enumerate(ocr["text"]):
      text = raw.strip()
      if not text:
         continue
      x, y, w, h = (
         ocr["left"][i], ocr["top"][i],
         ocr["width"][i], ocr["height"][i]
      )
      cx, cy = x + w // 2, y + h // 2
      for rect, (r, c) in cell_lookup.items():
         bx, by, bw, bh = rect
         if bx <= cx <= bx + bw and by <= cy <= by + bh:
            positions[text].append((r, c))
            break

   return {code: positions.get(code, []) for code in course_codes}


def annotate_courses_with_context(img: Any, course_codes: List[str]) -> Dict[str, List[Dict[str, Any]]]:
   """Combines course positions with day and venue info."""
   day_map, venue_map = extract_day_venue(img)
   pos_map = extract_course_positions(img, course_codes)
   result: Dict[str, List[Dict[str, Any]]] = {}
   for code, coords in pos_map.items():
      result[code] = []
      for r, c in coords:
         result[code].append({
            "row": r,
            "col": c,
            "day": day_map.get(r, ""),
            "venue": venue_map.get(r, "")
         })
   return result


def process_multiple_images(
   folder: Path,
   course_codes: List[str]
) -> Dict[str, Any]:
   """
   Processes all .png images in a folder, returns annotated info.
   """
   all_info: Dict[str, Any] = {}
   for img_path in folder.glob("*.png"):
      img = cv2.imread(str(img_path))
      if img is None:
         continue
      all_info[img_path.name] = annotate_courses_with_context(img, course_codes)

   return all_info


if __name__ == "__main__":
   img_file = Path("your_timetable_image.png")
   if not img_file.exists():
      print(f"Error: Image file '{img_file}' not found.")
      exit(1)

   img = cv2.imread(str(img_file))
   if img is None:
      print(f"Error: Failed to load image '{img_file}'.")
      exit(1)

   course_list = ["Monday", "PHY223", "CSC425"]
   single = annotate_courses_with_context(img, course_list)
   for code, info in single.items():
      print(f" - {code}: {info}")

   folder_path = Path("./timetables")
   if folder_path.is_dir():
      print(f"\nBatch results:")
      batch = process_multiple_images(folder_path, course_list)
      for fname, data in batch.items():
         print(f"{fname}:")
         for code, info in data.items():
            print(f" - {code}: {info}")