fixed zip and ner outputs

app.py

@@ -9,6 +9,7 @@ from datetime import datetime
 import numpy as np
 import os
 from gliner import GLiNER
+import json
 
 # Initialize GLiNER model
 gliner_model = GLiNER.from_pretrained("knowledgator/modern-gliner-bi-large-v1.0")
@@ -113,20 +114,18 @@ def run_example(image, model_id="Qwen/Qwen2-VL-7B-Instruct", run_ner=False, ner_
             threshold=0.3
         )
         
-        # Format the text with entity annotations
-        annotated_text = ocr_text
-        for entity in sorted(ner_results, key=lambda x: x["start"], reverse=True):
-            entity_text = entity["text"]
-            entity_label = entity["label"]
-            annotated_text = (
-                annotated_text[:entity["start"]] +
-                f"[{entity_text}]({entity_label})" +
-                annotated_text[entity["end"]:]
-            )
+        # Format entities for highlighting
+        entities = []
+        for entity in ner_results:
+            entities.append((
+                entity["start"],
+                entity["end"],
+                entity["label"]
+            ))
         
-        return f"OCR Text:\n{ocr_text}\n\nAnnotated Entities:\n{annotated_text}"
+        return [(ocr_text, entities)]
     
-    return ocr_text
+    return [(ocr_text, [])]
 
 css = """
   /* Overall app styling */
@@ -218,7 +217,7 @@ with gr.Blocks(css=css) as demo:
                 
                 submit_btn = gr.Button(value="Submit", elem_classes="submit-btn")
             with gr.Column(elem_classes="output-container"):
-                output_text = gr.Textbox(label="Output Text", elem_id="output")
+                output_text = gr.HighlightedText(label="Output Text", elem_id="output")
 
         # Show/hide NER labels based on checkbox
         ner_checkbox.change(
@@ -237,35 +236,57 @@ with gr.Blocks(css=css) as demo:
         filename = gr.Textbox(label="Save filename (without extension)", placeholder="Enter filename to save")
         download_btn = gr.Button("Download Image & Text", elem_classes="submit-btn")
 
-    def create_zip(image, text, fname):
-        if not image or not text or not fname:
+    def create_zip(image, text_data, fname):
+        if not isinstance(image, (Image.Image, np.ndarray)) or not text_data or not fname:
             return None
         
+        # Convert numpy array to PIL Image if needed
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        
         # Create a temporary directory
         with tempfile.TemporaryDirectory() as temp_dir:
-            # Save image with same extension
+            # Save image
             img_ext = image.format.lower() if hasattr(image, 'format') else 'png'
             img_path = os.path.join(temp_dir, f"{fname}.{img_ext}")
             image.save(img_path)
             
+            # Extract text and entities from the HighlightedText output
+            text, entities = text_data[0] if isinstance(text_data, list) else (text_data, [])
+            
             # Save text
             txt_path = os.path.join(temp_dir, f"{fname}.txt")
             with open(txt_path, 'w', encoding='utf-8') as f:
                 f.write(text)
             
+            # Create JSON with text, entities, and image info
+            json_data = {
+                "text": text,
+                "entities": [
+                    {
+                        "start": start,
+                        "end": end,
+                        "label": label,
+                        "text": text[start:end]
+                    }
+                    for start, end, label in entities
+                ],
+                "image_file": f"{fname}.{img_ext}"
+            }
+            
+            # Save JSON
+            json_path = os.path.join(temp_dir, f"{fname}.json")
+            with open(json_path, 'w', encoding='utf-8') as f:
+                json.dump(json_data, f, indent=2, ensure_ascii=False)
+            
             # Create zip file
             zip_path = os.path.join(temp_dir, f"{fname}.zip")
             with zipfile.ZipFile(zip_path, 'w') as zipf:
                 zipf.write(img_path, os.path.basename(img_path))
                 zipf.write(txt_path, os.path.basename(txt_path))
+                zipf.write(json_path, os.path.basename(json_path))
             
             return zip_path
 
-    download_btn.click(
-        create_zip,
-        inputs=[input_img, output_text, filename],
-        outputs=gr.File(label="Download")
-    )
-
 demo.queue(api_open=False)
 demo.launch(debug=True)