#!/usr/bin/env python3 """ Production script to train Amharic/English script detector """ import argparse import json import logging import sys from pathlib import Path from datetime import datetime import numpy as np import cv2 from sklearn.ensemble import RandomForestClassifier from sklearn.model_selection import train_test_split from sklearn.metrics import classification_report, confusion_matrix import joblib import mlflow from mlflow.sklearn import log_model import pandas as pd # Add parent directory to path sys.path.append(str(Path(__file__).parent.parent)) from app.utils.image_processing import ImageProcessor from app.analyzers.script_detector import ScriptDetector logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class ScriptDetectorTrainer: """Production-ready script detector trainer""" def __init__(self, experiment_name="script_detector"): self.experiment_name = experiment_name self.mlflow_tracking_uri = "http://localhost:5000" # MLflow tracking server mlflow.set_tracking_uri(self.mlflow_tracking_uri) mlflow.set_experiment(experiment_name) def extract_features(self, image_path: str, label: str) -> dict: """Extract features from image for script detection""" img = cv2.imread(str(image_path)) if img is None: return None # Convert to grayscale gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # Resize for consistency resized = cv2.resize(gray, (256, 256)) features = { 'mean_intensity': np.mean(resized), 'std_intensity': np.std(resized), 'skewness': float(pd.Series(resized.flatten()).skew()), 'kurtosis': float(pd.Series(resized.flatten()).kurtosis()), # Edge features 'edges_sobel': np.mean(cv2.Sobel(resized, cv2.CV_64F, 1, 1)), 'edges_canny': np.mean(cv2.Canny(resized, 100, 200)) / 255.0, # Texture features 'contrast': self._calculate_contrast(resized), 'homogeneity': self._calculate_homogeneity(resized), # Fourier transform features (text frequency) 'high_freq_energy': self._calculate_frequency_energy(resized), # Label 'label': 0 if label == 'eng' else 1 if label == 'amh' else 2, # mixed=2 'label_name': label } # Histogram features hist = cv2.calcHist([resized], [0], None, [16], [0, 256]).flatten() hist = hist / hist.sum() for i, val in enumerate(hist): features[f'hist_bin_{i}'] = float(val) return features def _calculate_contrast(self, image: np.ndarray) -> float: """Calculate image contrast""" min_val = np.min(image) max_val = np.max(image) return float((max_val - min_val) / (max_val + min_val + 1e-10)) def _calculate_homogeneity(self, image: np.ndarray) -> float: """Calculate image homogeneity""" from skimage.feature import graycomatrix, graycoprops try: glcm = graycomatrix(image.astype(np.uint8), [1], [0], symmetric=True, normed=True) return float(graycoprops(glcm, 'homogeneity')[0, 0]) except: return 0.5 def _calculate_frequency_energy(self, image: np.ndarray) -> float: """Calculate high-frequency energy in Fourier domain""" f = np.fft.fft2(image) fshift = np.fft.fftshift(f) magnitude_spectrum = 20 * np.log(np.abs(fshift) + 1) # High frequency energy (corners) h, w = magnitude_spectrum.shape center_h, center_w = h // 2, w // 2 corner_energy = np.sum(magnitude_spectrum[center_h-20:center_h+20, center_w-20:center_w+20]) total_energy = np.sum(magnitude_spectrum) return float(1 - (corner_energy / total_energy) if total_energy > 0 else 0) def prepare_dataset(self, data_dir: str) -> tuple: """Prepare dataset from directory structure""" data_dir = Path(data_dir) features_list = [] labels = [] # Expected structure: data_dir/{eng,amh,mixed}/*.png for script_type in ['eng', 'amh', 'mixed']: script_dir = data_dir / script_type if not script_dir.exists(): logger.warning(f"Directory not found: {script_dir}") continue image_files = list(script_dir.glob("*.png")) + list(script_dir.glob("*.jpg")) for img_path in image_files: features = self.extract_features(img_path, script_type) if features: features_list.append(features) labels.append(features['label']) logger.info(f"Loaded {len(image_files)} images for {script_type}") if not features_list: raise ValueError("No training data found") # Convert to DataFrame df = pd.DataFrame(features_list) # Prepare X and y feature_cols = [col for col in df.columns if not col.startswith(('label', 'hist_bin_'))] X = df[feature_cols].values y = df['label'].values logger.info(f"Dataset shape: {X.shape}, Labels: {np.unique(y, return_counts=True)}") return X, y, feature_cols def train(self, X_train, y_train, X_val, y_val, params: dict = None): """Train model with MLflow tracking""" if params is None: params = { 'n_estimators': 200, 'max_depth': 15, 'min_samples_split': 5, 'min_samples_leaf': 2, 'random_state': 42, 'n_jobs': -1 } with mlflow.start_run(): # Log parameters mlflow.log_params(params) # Train model model = RandomForestClassifier(**params) model.fit(X_train, y_train) # Evaluate train_score = model.score(X_train, y_train) val_score = model.score(X_val, y_val) y_pred = model.predict(X_val) report = classification_report(y_val, y_pred, output_dict=True) # Log metrics mlflow.log_metric("train_accuracy", train_score) mlflow.log_metric("val_accuracy", val_score) mlflow.log_metric("precision", report['weighted avg']['precision']) mlflow.log_metric("recall", report['weighted avg']['recall']) mlflow.log_metric("f1_score", report['weighted avg']['f1-score']) # Log confusion matrix cm = confusion_matrix(y_val, y_pred) cm_path = "confusion_matrix.png" self._plot_confusion_matrix(cm, ['eng', 'amh', 'mixed'], cm_path) mlflow.log_artifact(cm_path) # Log model model_info = mlflow.sklearn.log_model(model, "script_detector_model") logger.info(f"Training complete. Validation accuracy: {val_score:.4f}") return model, model_info def _plot_confusion_matrix(self, cm, classes, save_path): """Plot and save confusion matrix""" import matplotlib.pyplot as plt plt.figure(figsize=(8, 6)) plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues) plt.title('Confusion Matrix') plt.colorbar() tick_marks = np.arange(len(classes)) plt.xticks(tick_marks, classes, rotation=45) plt.yticks(tick_marks, classes) # Add text annotations thresh = cm.max() / 2. for i in range(cm.shape[0]): for j in range(cm.shape[1]): plt.text(j, i, format(cm[i, j], 'd'), horizontalalignment="center", color="white" if cm[i, j] > thresh else "black") plt.tight_layout() plt.ylabel('True label') plt.xlabel('Predicted label') plt.savefig(save_path, dpi=150, bbox_inches='tight') plt.close() def save_model(self, model, output_path: str, feature_cols: list): """Save model with metadata""" model_data = { 'model': model, 'feature_columns': feature_cols, 'version': '1.0.0', 'trained_at': datetime.now().isoformat(), 'classes': ['eng', 'amh', 'mixed'] } output_path = Path(output_path) output_path.parent.mkdir(parents=True, exist_ok=True) joblib.dump(model_data, output_path) logger.info(f"Model saved to {output_path}") # Save feature importance importance_df = pd.DataFrame({ 'feature': feature_cols, 'importance': model.feature_importances_ }).sort_values('importance', ascending=False) importance_path = output_path.parent / "feature_importance.csv" importance_df.to_csv(importance_path, index=False) logger.info(f"Feature importance saved to {importance_path}") def main(): parser = argparse.ArgumentParser(description="Train script detector for Amharic/English") parser.add_argument("--data-dir", required=True, help="Directory with training data") parser.add_argument("--output-dir", default="models/script_detector", help="Output directory") parser.add_argument("--test-size", type=float, default=0.2, help="Test set size") parser.add_argument("--random-state", type=int, default=42, help="Random seed") parser.add_argument("--mlflow", action="store_true", help="Enable MLflow tracking") args = parser.parse_args() # Start training trainer = ScriptDetectorTrainer() logger.info("Preparing dataset...") X, y, feature_cols = trainer.prepare_dataset(args.data_dir) # Split data X_train, X_val, y_train, y_val = train_test_split( X, y, test_size=args.test_size, random_state=args.random_state, stratify=y ) logger.info(f"Training set: {X_train.shape}, Validation set: {X_val.shape}") # Train model logger.info("Training model...") model, model_info = trainer.train(X_train, y_train, X_val, y_val) # Save model output_path = Path(args.output_dir) / "script_detector.joblib" trainer.save_model(model, output_path, feature_cols) logger.info("✅ Training completed successfully!") if __name__ == "__main__": main()