Update app.py

app.py

@@ -54,83 +54,154 @@ class DataManager:
             self.dataset = None
     
     def _load_segment_metadata(self):
-        """Load segment metadata from RAR file"""
+        """Load segment metadata from RAR file containing multiple JSON files"""
         RAR_FILE = "kiu_segment_metadata.rar"
-        JSON_FILE = os.path.join(SEGMENT_EXTRACT_DIR, "kiu_segment_metadata.json")
+        EXTRACT_DIR = SEGMENT_EXTRACT_DIR
         
         # Check if already extracted
-        if os.path.exists(JSON_FILE):
-            print(f"Loading cached segment metadata...")
-            with open(JSON_FILE, 'r') as f:
-                self.segment_data = json.load(f)
-            print(f"✅ Loaded {len(self.segment_data)} segment entries")
-            return
+        if os.path.exists(EXTRACT_DIR) and os.listdir(EXTRACT_DIR):
+            json_files = [f for f in os.listdir(EXTRACT_DIR) if f.endswith('.json')]
+            if json_files:
+                print(f"Loading cached segment metadata from {len(json_files)} files...")
+                self.segment_data = {}
+                
+                for json_file in json_files:
+                    filepath = os.path.join(EXTRACT_DIR, json_file)
+                    try:
+                        with open(filepath, 'r') as f:
+                            data = json.load(f)
+                            # Use filename as key (e.g., "KIU_00001_segments")
+                            key = json_file.replace('.json', '')
+                            self.segment_data[key] = data
+                    except Exception as e:
+                        print(f"⚠️ Error loading {json_file}: {e}")
+                
+                print(f"✅ Loaded {len(self.segment_data)} segment entries from cache")
+                return
         
         # Extract from RAR
         if not os.path.exists(RAR_FILE):
             print(f"❌ {RAR_FILE} not found in Space root")
+            print(f"Current directory: {os.getcwd()}")
+            print(f"Files available: {os.listdir('.')}")
             return
         
         print(f"Extracting {RAR_FILE}...")
+        
         try:
+            # Set rarfile to use unrar
+            rarfile.UNRAR_TOOL = "unrar"
+            
             with rarfile.RarFile(RAR_FILE) as rf:
-                # Find JSON file
-                json_in_rar = None
-                for file_info in rf.infolist():
-                    if file_info.filename.endswith('.json') and 'segment_metadata' in file_info.filename:
-                        json_in_rar = file_info.filename
-                        break
+                # List all files
+                all_files = [f.filename for f in rf.infolist()]
+                json_files = [f for f in all_files if f.endswith('.json')]
                 
-                if not json_in_rar:
-                    print("❌ No JSON file found in RAR")
+                print(f"Found {len(json_files)} JSON files in RAR")
+                
+                if not json_files:
+                    print(f"❌ No JSON files found in RAR")
+                    print(f"Files in RAR: {all_files[:10]}...")  # Show first 10
                     return
                 
-                # Extract
-                rf.extract(json_in_rar, SEGMENT_EXTRACT_DIR)
-                extracted = os.path.join(SEGMENT_EXTRACT_DIR, json_in_rar)
+                # Extract all JSON files
+                os.makedirs(EXTRACT_DIR, exist_ok=True)
+                
+                for json_file in json_files:
+                    try:
+                        rf.extract(json_file, EXTRACT_DIR)
+                    except Exception as e:
+                        print(f"⚠️ Error extracting {json_file}: {e}")
+                        # Try reading directly
+                        try:
+                            with rf.open(json_file) as f:
+                                content = f.read()
+                                # Get just the filename without path
+                                filename = os.path.basename(json_file)
+                                output_path = os.path.join(EXTRACT_DIR, filename)
+                                with open(output_path, 'wb') as out:
+                                    out.write(content)
+                        except Exception as e2:
+                            print(f"❌ Failed to extract {json_file}: {e2}")
+                            continue
                 
-                # Load
-                with open(extracted, 'r') as f:
-                    self.segment_data = json.load(f)
+                # Now load all extracted JSON files
+                self.segment_data = {}
+                extracted_files = [f for f in os.listdir(EXTRACT_DIR) if f.endswith('.json')]
                 
-                # Rename to standard location
-                if extracted != JSON_FILE:
-                    shutil.move(extracted, JSON_FILE)
+                print(f"Loading {len(extracted_files)} extracted JSON files...")
                 
-                print(f"✅ Extracted and loaded {len(self.segment_data)} entries")
+                for json_file in extracted_files:
+                    filepath = os.path.join(EXTRACT_DIR, json_file)
+                    try:
+                        with open(filepath, 'r') as f:
+                            data = json.load(f)
+                            # Use filename as key (e.g., "KIU_00001_segments")
+                            key = json_file.replace('.json', '')
+                            self.segment_data[key] = data
+                    except Exception as e:
+                        print(f"⚠️ Error loading {json_file}: {e}")
                 
+                print(f"✅ Extracted and loaded {len(self.segment_data)} segment entries")
+                
+        except rarfile.RarCannotExec as e:
+            print(f"❌ RAR tool not available: {e}")
+            print("⚠️ SOLUTION: Ensure packages.txt contains 'unrar'")
+            print("   Or extract manually and upload JSON files to:")
+            print(f"   {EXTRACT_DIR}/")
         except Exception as e:
-            print(f"❌ Error extracting RAR: {e}")
+            print(f"❌ Error with RAR file: {e}")
             import traceback
             traceback.print_exc()
     
     def _build_manifest(self):
-        """Build manifest of all line instances"""
+        """Build manifest of all line instances from multiple JSON files"""
         if not self.segment_data:
             print("❌ No segment data - cannot build manifest")
             return
         
-        print("Building manifest...")
+        print("Building manifest from segment metadata...")
         self.manifest = []
         
+        # self.segment_data is now a dict where:
+        # key = "KIU_00001_segments" (filename without .json)
+        # value = the JSON content for that KIU
+        
         for key, segment_info in self.segment_data.items():
             if not isinstance(segment_info, dict):
                 continue
             
-            # Extract KIU ID
-            kiu_id = segment_info.get('kiu_id')
-            if not kiu_id and 'KIU_' in key:
+            # Extract KIU ID from key (e.g., "KIU_00001_segments" -> "00001")
+            kiu_id = None
+            if 'KIU_' in key:
                 try:
-                    kiu_id = int(key.split('_')[1])
+                    # Extract the numeric part after KIU_
+                    parts = key.split('_')
+                    for part in parts:
+                        if part.isdigit() or (part.startswith('0') and part[1:].isdigit()):
+                            kiu_id = part
+                            break
                 except:
-                    continue
+                    pass
+            
+            # Also check if kiu_id is in the data itself
+            if not kiu_id and 'kiu_id' in segment_info:
+                kiu_id = str(segment_info['kiu_id'])
+            
+            if not kiu_id:
+                print(f"⚠️ Could not extract KIU ID from key: {key}")
+                continue
+            
+            # Ensure it's zero-padded to 5 digits
+            kiu_id = str(kiu_id).zfill(5)
             
             # Process line instances
             instances = segment_info.get('instances', [])
+            
             for instance in instances:
                 if instance.get('class') == 'Line':
                     self.manifest.append({
-                        'kiu_id': str(kiu_id).zfill(5),
+                        'kiu_id': kiu_id,
                         'instance_id': instance.get('instance_id'),
                         'crop_coords': instance.get('crop_coords', [0, 0, 100, 100]),
                         'direction': instance.get('direction', ''),
@@ -594,28 +665,60 @@ Skipped: {stats['skip']} | Unclear: {stats['unclear']} | Needs Reclass: {stats['
 app_state = AnnotationApp()
 
 # Build Gradio UI
-with gr.Blocks(title="⚡ Hieroglyph Annotation", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="⚡ Hieroglyph Annotation") as demo:
     
     gr.Markdown("# ⚡ Hieroglyph Direction Annotation Tool")
     
+    # Check if data is loaded
     stats = app_state.get_statistics()
+    data_loaded = len(app_state.data_mgr.manifest) > 0
     
-    gr.HTML(f"""
-    <div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); 
-                color: white; padding: 15px; border-radius: 10px; margin-bottom: 20px;">
-        <h3 style="margin: 0 0 10px 0;">📊 System Status</h3>
-        <div style="display: grid; grid-template-columns: repeat(3, 1fr); gap: 10px;">
-            <div><strong>Dataset:</strong> {HF_DATASET_REPO}</div>
-            <div><strong>Total Instances:</strong> {stats['total']:,}</div>
-            <div><strong>Remaining:</strong> {stats['remaining']:,}</div>
+    if not data_loaded:
+        gr.HTML(f"""
+        <div style="background: #fff3cd; padding: 20px; border-radius: 10px; margin-bottom: 20px; border-left: 5px solid #ffc107;">
+            <h3 style="margin-top: 0; color: #856404;">⚠️ No Segment Data Loaded</h3>
+            <p><strong>The segment metadata could not be loaded from the RAR file.</strong></p>
+            
+            <h4>Solution Options:</h4>
+            <ol>
+                <li><strong>Install unrar tool:</strong>
+                    <br>Add to your Space's Dockerfile or requirements:
+                    <pre style="background: #f8f9fa; padding: 10px; border-radius: 5px; margin: 10px 0;">RUN apt-get update && apt-get install -y unrar</pre>
+                </li>
+                <li><strong>Upload extracted JSON directly:</strong>
+                    <br>Extract <code>kiu_segment_metadata.json</code> locally and upload to:
+                    <pre style="background: #f8f9fa; padding: 10px; border-radius: 5px; margin: 10px 0;">{SEGMENT_EXTRACT_DIR}/kiu_segment_metadata.json</pre>
+                </li>
+                <li><strong>Use a different format:</strong>
+                    <br>Upload the JSON file as a regular file in your Space repo instead of RAR.
+                </li>
+            </ol>
+            
+            <p><strong>Current Status:</strong></p>
+            <ul>
+                <li>Dataset: {HF_DATASET_REPO} - {'✅ Loaded' if app_state.data_mgr.dataset else '❌ Failed'}</li>
+                <li>Segment Metadata: ❌ Not loaded</li>
+                <li>Manifest: {len(app_state.data_mgr.manifest)} instances</li>
+            </ul>
         </div>
-        <div style="margin-top: 10px; background: rgba(255,255,255,0.2); 
-                    padding: 8px; border-radius: 5px;">
-            <strong>Progress:</strong> {stats['progress_pct']:.1f}% 
-            ({stats['processed']:,} / {stats['total']:,} annotated)
+        """)
+    else:
+        gr.HTML(f"""
+        <div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); 
+                    color: white; padding: 15px; border-radius: 10px; margin-bottom: 20px;">
+            <h3 style="margin: 0 0 10px 0;">📊 System Status</h3>
+            <div style="display: grid; grid-template-columns: repeat(3, 1fr); gap: 10px;">
+                <div><strong>Dataset:</strong> {HF_DATASET_REPO}</div>
+                <div><strong>Total Instances:</strong> {stats['total']:,}</div>
+                <div><strong>Remaining:</strong> {stats['remaining']:,}</div>
+            </div>
+            <div style="margin-top: 10px; background: rgba(255,255,255,0.2); 
+                        padding: 8px; border-radius: 5px;">
+                <strong>Progress:</strong> {stats['progress_pct']:.1f}% 
+                ({stats['processed']:,} / {stats['total']:,} annotated)
+            </div>
         </div>
-    </div>
-    """)
+        """)
     
     with gr.Row():
         with gr.Column(scale=2):
@@ -713,4 +816,4 @@ if __name__ == "__main__":
     print(f"Progress: {stats['progress_pct']:.1f}%")
     print("="*80 + "\n")
     
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860, theme=gr.themes.Soft())