Update app.py

app.py

@@ -1,96 +1,247 @@
-# app_working.py
+# app_fixed.py
 import gradio as gr
-import os
 import json
-from datasets import load_dataset
+import os
+import pickle
 from PIL import Image
+import tempfile
+import time
+from pathlib import Path
+import hashlib
 import io
+import shutil
+import rarfile
 
-HF_TOKEN = os.getenv("HF_TOKEN")
-DATASET_REPO = "alyex/karnak-data-app"
+# Hugging Face imports
+from datasets import load_dataset
 
-# Simple state
-current_index = 0
-manifest = []
+# ========== CONFIGURATION ==========
+HF_DATASET_REPO = "alyex/karnak-data-app"
 
-def load_test_manifest():
-    """Create a test manifest"""
-    return [
-        {"kiu_id": "00001", "instance_id": 0, "crop_coords": [0, 0, 100, 100]},
-        {"kiu_id": "00002", "instance_id": 0, "crop_coords": [0, 0, 100, 100]},
-    ]
+# Directories
+HF_CACHE_DIR = os.path.join(tempfile.gettempdir(), "karnak_cache")
+HF_ANNOTATIONS_DIR = "/data/annotations"
+HF_HARD_NEGATIVES_DIR = "/data/hard_negatives"
+SEGMENT_EXTRACT_DIR = "/data/kiu_segment_metadata"
 
-def get_current_data():
-    """Get data for current item"""
-    global current_index, manifest
-    
-    if not manifest:
-        manifest = load_test_manifest()
-    
-    if current_index >= len(manifest):
-        return None, "No items", "No data", ""
+os.makedirs(HF_CACHE_DIR, exist_ok=True)
+os.makedirs(HF_ANNOTATIONS_DIR, exist_ok=True)
+os.makedirs(HF_HARD_NEGATIVES_DIR, exist_ok=True)
+os.makedirs(SEGMENT_EXTRACT_DIR, exist_ok=True)
+
+MANIFEST_FILE = os.path.join(HF_CACHE_DIR, "manifest.pkl")
+
+class AnnotationState:
+    def __init__(self):
+        self.manifest = []
+        self.current_index = 0
+        self.history = []
+        self.karnak_dataset = None
+        self.segment_data = {}
+        self.kiu_lookup = {}  # Fast lookup dictionary
+        
+        # Load datasets
+        self.load_datasets()
+        # Load manifest
+        self.load_manifest()
     
-    item = manifest[current_index]
-    kiu_id = item["kiu_id"]
+    def load_datasets(self):
+        """Load datasets with fast lookup"""
+        print("Loading Karnak dataset...")
+        try:
+            self.karnak_dataset = load_dataset(HF_DATASET_REPO, streaming=True)
+            print(f"✅ Karnak dataset loaded")
+            
+            # Build fast lookup for first N items
+            self.build_kiu_lookup(2000)  # Lookup first 2000 KIUs
+            
+        except Exception as e:
+            print(f"❌ Error loading dataset: {e}")
+            self.karnak_dataset = None
+        
+        print("Loading segment metadata from local...")
+        self.load_segment_metadata_local()
     
-    # Try to fetch from dataset
-    try:
-        dataset = load_dataset(DATASET_REPO, streaming=True, token=HF_TOKEN)
-        
-        for data_item in dataset["train"].take(100):
-            if data_item["kiu_id"] == kiu_id:
-                html = data_item.get("html", "No HTML")
-                image = None
-                
-                if data_item.get("image"):
-                    try:
-                        image = Image.open(io.BytesIO(data_item["image"]))
-                        # Apply crop
-                        x1, y1, x2, y2 = item["crop_coords"]
-                        image = image.crop((x1, y1, x2, y2))
-                    except:
-                        pass
+    def build_kiu_lookup(self, limit=2000):
+        """Build fast lookup dictionary"""
+        if not self.karnak_dataset:
+            return
+        
+        print(f"Building KIU lookup table (first {limit} items)...")
+        self.kiu_lookup = {}
+        
+        count = 0
+        try:
+            for item in self.karnak_dataset["train"]:
+                kiu_id = str(item["kiu_id"]).strip()
+                # Normalize: remove leading zeros
+                kiu_id_norm = kiu_id.lstrip('0') or '0'
                 
-                html_display = f"<div style='padding:10px'>{html[:1000]}...</div>"
-                info = f"KIU {kiu_id} | Instance {item['instance_id']}"
+                # Store in lookup
+                self.kiu_lookup[kiu_id_norm] = item
                 
-                return image, html_display, info, ""
+                count += 1
+                if count >= limit:
+                    break
+            
+            print(f"✅ Built lookup for {count} KIUs")
+            
+        except Exception as e:
+            print(f"Lookup building stopped: {e}")
     
-    except Exception as e:
-        return None, f"Error: {str(e)}", f"KIU {kiu_id}", ""
-    
-    return None, "Not found in dataset", f"KIU {kiu_id}", ""
-
-with gr.Blocks(title="⚡ Working App") as demo:
-    gr.Markdown("# ⚡ Hieroglyph Annotation")
-    
-    image_display = gr.Image(label="Line Instance")
-    html_display = gr.HTML(label="Reference")
-    info_display = gr.Markdown()
-    line_input = gr.Textbox(label="Line Number")
-    
-    next_btn = gr.Button("Next")
-    prev_btn = gr.Button("Previous")
-    
-    def update_display():
-        return get_current_data()
+    def load_segment_metadata_local(self):
+        """Load from local RAR file"""
+        RAR_FILENAME = "kiu_segment_metadata.rar"
+        
+        if not os.path.exists(RAR_FILENAME):
+            print(f"❌ {RAR_FILENAME} not found")
+            return
+        
+        print(f"✅ Found {RAR_FILENAME}")
+        
+        # Check cache
+        json_path = os.path.join(SEGMENT_EXTRACT_DIR, "kiu_segment_metadata.json")
+        if os.path.exists(json_path):
+            print(f"Loading cached metadata...")
+            try:
+                with open(json_path, 'r') as f:
+                    self.segment_data = json.load(f)
+                print(f"✅ Loaded {len(self.segment_data)} entries")
+                return
+            except:
+                pass
+        
+        # Extract from RAR
+        try:
+            print(f"Extracting {RAR_FILENAME}...")
+            with rarfile.RarFile(RAR_FILENAME) as rf:
+                # Find and extract JSON
+                for file_info in rf.infolist():
+                    if file_info.filename.endswith('.json'):
+                        rf.extract(file_info, SEGMENT_EXTRACT_DIR)
+                        extracted_path = os.path.join(SEGMENT_EXTRACT_DIR, file_info.filename)
+                        
+                        # Load data
+                        with open(extracted_path, 'r') as f:
+                            self.segment_data = json.load(f)
+                        
+                        print(f"✅ Loaded {len(self.segment_data)} entries")
+                        
+                        # Cache as standard name
+                        if extracted_path != json_path:
+                            shutil.move(extracted_path, json_path)
+                        
+                        break
+        
+        except Exception as e:
+            print(f"❌ Extraction error: {e}")
     
-    def next_item():
-        global current_index
-        if current_index < len(manifest) - 1:
-            current_index += 1
-        return get_current_data()
+    def load_manifest(self):
+        """Build manifest, filtering for KIUs we have in dataset"""
+        if not self.segment_data:
+            print("❌ No segment data")
+            self.manifest = []
+            return
+        
+        print("Building manifest (checking for available KIUs)...")
+        self.manifest = []
+        
+        available_count = 0
+        unavailable_count = 0
+        
+        for key, segment_info in self.segment_data.items():
+            if not isinstance(segment_info, dict):
+                continue
+            
+            # Get KIU ID from segment metadata
+            kiu_id = str(segment_info.get('kiu_id', ''))
+            kiu_id_norm = kiu_id.lstrip('0') or '0'
+            
+            # Check if we have this KIU in our lookup
+            if kiu_id_norm in self.kiu_lookup:
+                available_count += 1
+                instances = segment_info.get('instances', [])
+                
+                for instance in instances:
+                    if instance.get('class') == 'Line':
+                        manifest_item = {
+                            'kiu_id': kiu_id.zfill(5),  # Store with leading zeros
+                            'kiu_id_norm': kiu_id_norm,  # Normalized for lookup
+                            'instance_id': instance.get('instance_id'),
+                            'crop_coords': instance.get('crop_coords', [0, 0, 100, 100]),
+                            'direction': instance.get('direction', ''),
+                            'annotated': instance.get('annotated', False),
+                            'line_number': instance.get('line_number'),
+                            'segment_key': key
+                        }
+                        self.manifest.append(manifest_item)
+            else:
+                unavailable_count += 1
+        
+        print(f"✅ Manifest: {len(self.manifest)} instances from {available_count} available KIUs")
+        print(f"   Skipped: {unavailable_count} KIUs not in dataset")
+        
+        if self.manifest:
+            # Save to cache
+            with open(MANIFEST_FILE, 'wb') as f:
+                pickle.dump(self.manifest, f)
+            
+            self.current_index = self.find_first_unannotated()
     
-    def prev_item():
-        global current_index
-        if current_index > 0:
-            current_index -= 1
-        return get_current_data()
+    def get_kiu_data(self, kiu_id):
+        """Fast lookup using dictionary"""
+        if not self.kiu_lookup:
+            return None, None
+        
+        # Normalize KIU ID
+        kiu_id_norm = str(kiu_id).strip().lstrip('0') or '0'
+        
+        if kiu_id_norm in self.kiu_lookup:
+            item = self.kiu_lookup[kiu_id_norm]
+            
+            # Convert image
+            image = None
+            if item.get("image"):
+                try:
+                    image = Image.open(io.BytesIO(item["image"]))
+                except Exception as e:
+                    print(f"Image conversion error for KIU {kiu_id}: {e}")
+            
+            return item.get("html", ""), image
+        
+        # Fallback: search in dataset
+        return self.get_kiu_data_slow(kiu_id)
     
-    # Initial load
-    demo.load(update_display, outputs=[image_display, html_display, info_display, line_input])
+    def get_kiu_data_slow(self, kiu_id):
+        """Fallback: search through dataset"""
+        if not self.karnak_dataset:
+            return None, None
+        
+        kiu_id_norm = str(kiu_id).strip().lstrip('0') or '0'
+        
+        try:
+            for item in self.karnak_dataset["train"]:
+                item_id = str(item["kiu_id"]).strip().lstrip('0') or '0'
+                if item_id == kiu_id_norm:
+                    image = None
+                    if item.get("image"):
+                        try:
+                            image = Image.open(io.BytesIO(item["image"]))
+                        except:
+                            pass
+                    return item.get("html", ""), image
+        except:
+            pass
+        
+        return None, None
     
-    next_btn.click(next_item, outputs=[image_display, html_display, info_display, line_input])
-    prev_btn.click(prev_item, outputs=[image_display, html_display, info_display, line_input])
+    # [Keep all your other methods: needs_annotation, find_first_unannotated, etc.]
+    # ...
+
+# [Keep the rest of your app code: load_current_instance, navigation functions, Gradio interface]
+# ...
+
+# Initialize state
+state = AnnotationState()
 
-demo.launch(server_name="0.0.0.0", server_port=7860)
+# Rest of your app code remains the same...