updated the annotation library

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
 import os
 import json
-from huggingface_hub import hf_hub_download, list_repo_files
+from huggingface_hub import hf_hub_download, list_repo_files, upload_file, HfApi
 import logging
+import tempfile
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -14,8 +15,8 @@ ANNOTATION_CATEGORIES = {
     "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", "six", "wicket", "runs (1,2,3)"],
-    "Shot direction": ["long on", "long off", "cover", "point", "midwicket", "square leg", "third man", "fine leg"],
+    "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"]
 }
 
@@ -28,6 +29,7 @@ class VideoAnnotator:
         self.video_files = []
         self.current_video_idx = 0
         self.annotations = {}
+        self.hf_token = os.environ.get("HF_TOKEN")  # Get token from environment variable
         
     def load_videos_from_hf(self):
         try:
@@ -61,20 +63,72 @@ class VideoAnnotator:
             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])
         annotation_file = f"{video_name}_annotations.json"
         
         logger.info(f"Saving annotations for {video_name} to {annotation_file}")
         
         try:
-            with open(annotation_file, 'w') as f:
+            # Save locally in a temporary file
+            with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
                 json.dump(annotations_dict, f, indent=4)
-            logger.info("Annotations saved successfully")
-            return f"Annotations saved for {video_name}"
+                temp_path = f.name
+            
+            logger.info("Annotations saved locally")
+            
+            # Upload to Hugging Face repo
+            if self.hf_token:
+                logger.info(f"Uploading annotations to Hugging Face repo: {HF_REPO_ID}")
+                api = HfApi()
+                api.upload_file(
+                    path_or_fileobj=temp_path,
+                    path_in_repo=f"annotations/{annotation_file}",
+                    repo_id=HF_REPO_ID,
+                    repo_type=HF_REPO_TYPE,
+                    token=self.hf_token
+                )
+                os.unlink(temp_path)  # Clean up the temporary file
+                return f"Annotations saved for {video_name} and uploaded to Hugging Face repo"
+            else:
+                logger.warning("HF_TOKEN not found. Annotations saved locally only.")
+                # Copy from temp to current directory for local access
+                with open(annotation_file, 'w') as f:
+                    with open(temp_path, 'r') as temp_f:
+                        f.write(temp_f.read())
+                os.unlink(temp_path)  # Clean up the temporary file
+                return f"Annotations saved locally for {video_name} (no HF upload)"
         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 repo"""
+        if not self.video_files:
+            return None
+            
+        video_name = os.path.basename(self.video_files[self.current_video_idx])
+        annotation_path = f"annotations/{video_name}_annotations.json"
+        
+        try:
+            # Check if annotation exists in the repo
+            all_files = list_repo_files(HF_REPO_ID, repo_type=HF_REPO_TYPE)
+            if annotation_path in all_files:
+                local_path = hf_hub_download(
+                    repo_id=HF_REPO_ID,
+                    filename=annotation_path,
+                    repo_type=HF_REPO_TYPE
+                )
+                with open(local_path, 'r') as f:
+                    return json.load(f)
+            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:
@@ -159,6 +213,13 @@ def create_interface():
         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(