fixed ner and text output

app.py

@@ -257,11 +257,7 @@ with gr.Blocks(css=css) as demo:
 
     def create_zip(image, text_data, fname):
         # Validate inputs
-        if not fname:
-            return None
-        if not text_data:
-            return None
-        if not isinstance(image, (Image.Image, np.ndarray)):
+        if not fname or not text_data or not isinstance(image, (Image.Image, np.ndarray)):
             return None
         
         try:
@@ -271,27 +267,35 @@ with gr.Blocks(css=css) as demo:
             
             # Create a temporary directory
             with tempfile.TemporaryDirectory() as temp_dir:
-                # Save image with default png extension
+                # Save image
                 img_path = os.path.join(temp_dir, f"{fname}.png")
                 image.save(img_path)
                 
-                # Extract text content
-                text = ""
+                # Extract text and entities from text_data
+                full_text = ""
                 entities = []
-                if isinstance(text_data, list) and len(text_data) > 0:
-                    if isinstance(text_data[0], tuple):
-                        text = text_data[0][0]  # Get text from first tuple
-                    else:
-                        text = str(text_data[0])  # Convert to string if not tuple
+                current_pos = 0
+                
+                # Process the highlighted text data
+                for segment, label in text_data:
+                    full_text += segment
+                    if label:  # If this segment has a label (is an entity)
+                        entities.append({
+                            "text": segment,
+                            "label": label,
+                            "start": current_pos,
+                            "end": current_pos + len(segment)
+                        })
+                    current_pos += len(segment)
                 
                 # 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)
+                    f.write(full_text)
                 
                 # Create JSON with text and entities
                 json_data = {
-                    "text": text,
+                    "text": full_text,
                     "entities": entities,
                     "image_file": f"{fname}.png"
                 }
@@ -301,7 +305,7 @@ with gr.Blocks(css=css) as demo:
                 with open(json_path, 'w', encoding='utf-8') as f:
                     json.dump(json_data, f, indent=2, ensure_ascii=False)
                 
-                # Create zip file in a location that will persist
+                # Create zip file
                 output_dir = "downloads"
                 os.makedirs(output_dir, exist_ok=True)
                 zip_path = os.path.join(output_dir, f"{fname}.zip")
@@ -311,7 +315,7 @@ with gr.Blocks(css=css) as demo:
                     zipf.write(txt_path, os.path.basename(txt_path))
                     zipf.write(json_path, os.path.basename(json_path))
                 
-                return zip_path  # Return just the path to the zip file
+                return zip_path
 
         except Exception as e:
             print(f"Error creating zip: {str(e)}")