import json import pandas as pd from datetime import datetime from typing import List, Dict, Any, Optional from dataclasses import dataclass from gc_storage import get_blob_names_with_prefix, download_blob_stream from settings import BUCKET_NAME, GOOGLE_CLOUD_PROJECT, ANNOTATIONS_FOLDER @dataclass class AnnotationMetrics: """Data class to hold annotation metrics""" total_annotations: int prediction_accuracy: float avg_lead_time: float annotations_per_user: Dict[str, int] label_distribution: Dict[str, int] throughput_by_date: Dict[str, int] correction_rate: float class DataProcessor: """Process annotation data from Google Cloud Storage""" def __init__(self, bucket_name: Optional[str] = None, project_id: Optional[str] = None): self.bucket_name = bucket_name or BUCKET_NAME self.project_id = project_id or GOOGLE_CLOUD_PROJECT # Validate required environment variables if not self.bucket_name: raise ValueError("BUCKET_NAME environment variable is required") if not self.project_id: raise ValueError("GOOGLE_CLOUD_PROJECT environment variable is required") def load_annotations_from_prefix(self, prefix: Optional[str] = None) -> List[Dict[str, Any]]: """Load all annotation files from a bucket prefix""" if prefix is None: prefix = ANNOTATIONS_FOLDER if not prefix: raise ValueError("ANNOTATIONS_FOLDER environment variable is required") # Clean up the prefix to avoid double slashes prefix = prefix.strip('/') if prefix: prefix = prefix + '/' annotations = [] try: # Get list of blobs with the specified prefix blob_names = get_blob_names_with_prefix(self.bucket_name, prefix) for blob_name in blob_names: try: # Download and parse JSON content content = download_blob_stream(self.bucket_name, blob_name) if content: annotation = json.loads(content.decode('utf-8')) annotations.append(annotation) except (json.JSONDecodeError, UnicodeDecodeError) as e: print(f"Error parsing {blob_name}: {e}") continue except Exception as e: print(f"Error loading annotations from prefix {prefix}: {e}") return annotations def calculate_prediction_accuracy(self, annotations: List[Dict[str, Any]]) -> float: """Calculate the ratio of unchanged predictions vs corrected predictions""" if not annotations: return 0.0 total_predictions = 0 correct_predictions = 0 for annotation in annotations: results = annotation.get('result', []) for result in results: if result.get('type') == 'labels': total_predictions += 1 if result.get('origin') == 'prediction': correct_predictions += 1 return correct_predictions / total_predictions if total_predictions > 0 else 0.0 def calculate_throughput_by_date(self, annotations: List[Dict[str, Any]]) -> Dict[str, int]: """Calculate number of annotations completed per date""" throughput = {} for annotation in annotations: created_at = annotation.get('created_at', '') if created_at: try: date = datetime.fromisoformat(created_at.replace('Z', '+00:00')).date() date_str = date.strftime('%Y-%m-%d') throughput[date_str] = throughput.get(date_str, 0) + 1 except ValueError: continue return throughput def calculate_user_productivity(self, annotations: List[Dict[str, Any]]) -> Dict[str, int]: """Calculate annotations per user""" user_counts = {} for annotation in annotations: completed_by = annotation.get('completed_by', {}) user_email = completed_by.get('email', 'Unknown') user_counts[user_email] = user_counts.get(user_email, 0) + 1 return user_counts def calculate_label_distribution(self, annotations: List[Dict[str, Any]]) -> Dict[str, int]: """Calculate frequency of each label type""" label_counts = {} for annotation in annotations: results = annotation.get('result', []) for result in results: if result.get('type') == 'labels': labels = result.get('value', {}).get('labels', []) for label in labels: label_counts[label] = label_counts.get(label, 0) + 1 return label_counts def calculate_avg_lead_time(self, annotations: List[Dict[str, Any]]) -> float: """Calculate average lead time for annotations""" lead_times = [] for annotation in annotations: lead_time = annotation.get('lead_time') if lead_time is not None: lead_times.append(float(lead_time)) return sum(lead_times) / len(lead_times) if lead_times else 0.0 def calculate_correction_rate(self, annotations: List[Dict[str, Any]]) -> float: """Calculate the rate of predictions that were corrected""" if not annotations: return 0.0 total_predictions = 0 corrected_predictions = 0 for annotation in annotations: results = annotation.get('result', []) for result in results: if result.get('type') == 'labels': total_predictions += 1 if result.get('origin') == 'prediction-changed': corrected_predictions += 1 return corrected_predictions / total_predictions if total_predictions > 0 else 0.0 def process_annotations(self, prefix: Optional[str] = None) -> AnnotationMetrics: """Process all annotations from a prefix and return metrics""" annotations = self.load_annotations_from_prefix(prefix) if not annotations: return AnnotationMetrics( total_annotations=0, prediction_accuracy=0.0, avg_lead_time=0.0, annotations_per_user={}, label_distribution={}, throughput_by_date={}, correction_rate=0.0 ) return AnnotationMetrics( total_annotations=len(annotations), prediction_accuracy=self.calculate_prediction_accuracy(annotations), avg_lead_time=self.calculate_avg_lead_time(annotations), annotations_per_user=self.calculate_user_productivity(annotations), label_distribution=self.calculate_label_distribution(annotations), throughput_by_date=self.calculate_throughput_by_date(annotations), correction_rate=self.calculate_correction_rate(annotations) ) def get_annotations_dataframe(self, prefix: Optional[str] = None) -> pd.DataFrame: """Convert annotations to pandas DataFrame for advanced analysis""" annotations = self.load_annotations_from_prefix(prefix) if not annotations: return pd.DataFrame() # Flatten annotation data for DataFrame flattened_data = [] for annotation in annotations: base_data = { 'id': annotation.get('id'), 'created_at': annotation.get('created_at'), 'updated_at': annotation.get('updated_at'), 'lead_time': annotation.get('lead_time'), 'user_email': annotation.get('completed_by', {}).get('email', 'Unknown'), 'task_id': annotation.get('task', {}).get('id'), 'project_id': annotation.get('project') } # Add result data results = annotation.get('result', []) for result in results: if result.get('type') == 'labels': row_data = base_data.copy() row_data.update({ 'result_id': result.get('id'), 'origin': result.get('origin'), 'labels': ', '.join(result.get('value', {}).get('labels', [])), 'start_time': result.get('value', {}).get('start'), 'end_time': result.get('value', {}).get('end'), 'channel': result.get('value', {}).get('channel') }) flattened_data.append(row_data) df = pd.DataFrame(flattened_data) # Convert datetime columns if 'created_at' in df.columns: df['created_at'] = pd.to_datetime(df['created_at']) if 'updated_at' in df.columns: df['updated_at'] = pd.to_datetime(df['updated_at']) return df