app.py

import gradio as gr
import os
import json
from huggingface_hub import hf_hub_download, list_repo_files, upload_file, HfApi
from datasets import load_dataset, Dataset
import logging
import tempfile

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Cricket annotation categories
ANNOTATION_CATEGORIES = {
    "Bowler's Run Up": ["Fast", "Slow"],
    "Delivery Type": ["Yorker", "Bouncer", "Length Ball", "Slower ball", "Googly", "Arm Ball", "Other"],
    "Ball's trajectory": ["In Swing", "Out Swing", "Off spin", "Leg spin"],
    "Shot Played": ["Cover Drive", "Straight Drive", "On Drive", "Pull", "Square Cut", "Defensive Block"],
    "Outcome of the shot": ["Four (4)", "Six (6)", "Wicket", "Single (1)", "Double (2)", "Triple (3)", "Dot (0)"],
    "Shot direction": ["Long On", "Long Off", "Cover", "Point", "Midwicket", "Square Leg", "Third Man", "Fine Leg"],
    "Fielder's Action": ["Catch taken", "Catch dropped", "Misfield", "Run-out attempt", "Fielder fields"]
}

HF_REPO_ID = "srrthk/CricBench" 
HF_REPO_TYPE = "dataset"  

class VideoAnnotator:
    def __init__(self):
        self.video_files = []
        self.current_video_idx = 0
        self.annotations = {}
        self.hf_token = os.environ.get("HF_TOKEN")
        self.dataset = None
        
    def load_videos_from_hf(self):
        try:
            logger.info(f"Loading dataset from HuggingFace: {HF_REPO_ID}")
            self.dataset = load_dataset(HF_REPO_ID, token=self.hf_token)
            
            # Get the split (usually 'train')
            split = list(self.dataset.keys())[0]
            self.dataset_split = self.dataset[split]
            
            # Get all video files from the dataset
            self.video_files = [item['video'] if 'video' in item else item['path'] 
                               for item in self.dataset_split]
            
            logger.info(f"Found {len(self.video_files)} video files")
            return len(self.video_files) > 0
        except Exception as e:
            logger.error(f"Error accessing HuggingFace dataset: {e}")
            return False
    
    def get_current_video(self):
        if not self.video_files:
            logger.warning("No video files available")
            return None
        
        video_path = self.video_files[self.current_video_idx]
        logger.info(f"Loading video: {video_path}")
        
        try:
            local_path = hf_hub_download(
                repo_id=HF_REPO_ID,
                filename=video_path,
                repo_type=HF_REPO_TYPE
            )
            logger.info(f"Video downloaded to: {local_path}")
            return local_path
        except Exception as e:
            logger.error(f"Error downloading video: {e}")
            return None
    
    def save_annotation(self, annotations_dict):
        if not annotations_dict:
            logger.warning("No annotations to save")
            return "No annotations to save"
            
        video_name = os.path.basename(self.video_files[self.current_video_idx])
        
        logger.info(f"Saving annotations for {video_name}")
        
        try:
            # Update the dataset with the new annotations
            if self.dataset is not None:
                # Get the split name (e.g., 'train')
                split = list(self.dataset.keys())[0]
                
                # Create a copy of the dataset to modify
                updated_dataset = self.dataset[split].to_pandas()
                
                # Convert annotations to JSON string
                annotation_json = json.dumps(annotations_dict)
                
                # Update the annotations column for the current video
                updated_dataset.loc[self.current_video_idx, 'annotations'] = annotation_json
                
                # Convert back to Hugging Face dataset
                new_dataset = Dataset.from_pandas(updated_dataset)
                
                # Push updated dataset to Hugging Face Hub
                if self.hf_token:
                    logger.info(f"Uploading updated dataset to Hugging Face: {HF_REPO_ID}")
                    new_dataset.push_to_hub(
                        HF_REPO_ID, 
                        split=split,
                        token=self.hf_token
                    )
                    # Update our local copy
                    self.dataset[split] = new_dataset
                    return f"Annotations saved for {video_name} and uploaded to Hugging Face dataset"
                else:
                    logger.warning("HF_TOKEN not found. Dataset updated locally only.")
                    self.dataset[split] = new_dataset
                    return f"Annotations saved locally for {video_name} (no HF upload)"
            else:
                logger.error("Dataset not loaded, cannot save annotations")
                return "Error: Dataset not loaded"
        except Exception as e:
            logger.error(f"Error saving annotations: {e}")
            return f"Error saving: {str(e)}"
    
    def load_existing_annotation(self):
        """Try to load existing annotation for the current video from the dataset"""
        if not self.dataset or not self.video_files:
            return None
            
        try:
            # Get the split name (e.g., 'train')
            split = list(self.dataset.keys())[0]
            
            # Check if the current item has annotations
            if 'annotations' in self.dataset[split][self.current_video_idx]:
                annotation_str = self.dataset[split][self.current_video_idx]['annotations']
                if annotation_str:
                    return json.loads(annotation_str)
            return None
        except Exception as e:
            logger.error(f"Error loading existing annotation: {e}")
            return None
    
    def next_video(self, *current_annotations):
        # Save current annotations before moving to next video
        if self.video_files:
            annotations_dict = {}
            for i, category in enumerate(ANNOTATION_CATEGORIES.keys()):
                if current_annotations[i]:
                    annotations_dict[category] = current_annotations[i]
            
            if annotations_dict:
                self.save_annotation(annotations_dict)
        
        # Move to next video
        if self.current_video_idx < len(self.video_files) - 1:
            self.current_video_idx += 1
            logger.info(f"Moving to next video (index: {self.current_video_idx})")
            return self.get_current_video(), *[None] * len(ANNOTATION_CATEGORIES)
        else:
            logger.info("Already at the last video")
            return self.get_current_video(), *[None] * len(ANNOTATION_CATEGORIES)
    
    def prev_video(self, *current_annotations):
        # Save current annotations before moving to previous video
        if self.video_files:
            annotations_dict = {}
            for i, category in enumerate(ANNOTATION_CATEGORIES.keys()):
                if current_annotations[i]:
                    annotations_dict[category] = current_annotations[i]
            
            if annotations_dict:
                self.save_annotation(annotations_dict)
        
        # Move to previous video
        if self.current_video_idx > 0:
            self.current_video_idx -= 1
            logger.info(f"Moving to previous video (index: {self.current_video_idx})")
            return self.get_current_video(), *[None] * len(ANNOTATION_CATEGORIES)
        else:
            logger.info("Already at the first video")
            return self.get_current_video(), *[None] * len(ANNOTATION_CATEGORIES)

def create_interface():
    annotator = VideoAnnotator()
    success = annotator.load_videos_from_hf()
    
    if not success:
        logger.error("Failed to load videos. Using demo mode with sample video.")
        # In real app, you might want to provide a sample video or show an error
    
    with gr.Blocks() as demo:
        gr.Markdown("# Cricket Video Annotation Tool")
        
        with gr.Row():
            video_player = gr.Video(label="Current Video")
        
        annotation_components = []
        
        with gr.Row():
            with gr.Column():
                for category, options in list(ANNOTATION_CATEGORIES.items())[:4]:
                    radio = gr.Radio(
                        choices=options,
                        label=category,
                        info=f"Select {category}"
                    )
                    annotation_components.append(radio)
            
            with gr.Column():
                for category, options in list(ANNOTATION_CATEGORIES.items())[4:]:
                    radio = gr.Radio(
                        choices=options,
                        label=category,
                        info=f"Select {category}"
                    )
                    annotation_components.append(radio)
        
        with gr.Row():
            prev_btn = gr.Button("Previous Video")
            save_btn = gr.Button("Save Annotations", variant="primary")
            next_btn = gr.Button("Next Video")
            
        # Initialize with first video
        current_video = annotator.get_current_video()
        if current_video:
            video_player.value = current_video
            
            # Try to load existing annotations
            existing_annotations = annotator.load_existing_annotation()
            if existing_annotations:
                for i, category in enumerate(ANNOTATION_CATEGORIES.keys()):
                    if category in existing_annotations:
                        annotation_components[i].value = existing_annotations[category]
        
        # Event handlers
        save_btn.click(
            fn=annotator.save_annotation,
            inputs=[gr.Group(annotation_components)],
            outputs=gr.Textbox(label="Status")
        )
        
        next_btn.click(
            fn=annotator.next_video,
            inputs=annotation_components,
            outputs=[video_player] + annotation_components
        )
        
        prev_btn.click(
            fn=annotator.prev_video,
            inputs=annotation_components,
            outputs=[video_player] + annotation_components
        )
        
    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch()
    # Add a local video for testing if no videos are loaded from Hugging Face