Update app.py

app.py

@@ -25,12 +25,6 @@ DEFAULT_NER_LABELS = "person, organization, location, date, event"
 
 # }
 
-class TextWithMetadata(list):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args)
-        self.original_text = kwargs.get('original_text', '')
-        self.entities = kwargs.get('entities', [])
-
 def array_to_image_path(image_array):
     # Convert numpy array to PIL Image
     img = Image.fromarray(np.uint8(image_array))
@@ -80,92 +74,104 @@ prompt_suffix = "<|end|>\n"
 
 @spaces.GPU
 def run_example(image, model_id="Qwen/Qwen2.5-VL-7B-Instruct", run_ner=False, ner_labels=DEFAULT_NER_LABELS):
-    # First get the OCR text
-    text_input = "Convert the image to text."
-    image_path = array_to_image_path(image)
-    
-    model = models[model_id]
-    processor = processors[model_id]
-    
-    prompt = f"{user_prompt}<|image_1|>\n{text_input}{prompt_suffix}{assistant_prompt}"
-    image = Image.fromarray(image).convert("RGB")
-    messages = [
-    {
-            "role": "user",
-            "content": [
-                {
-                    "type": "image",
-                    "image": image_path,
-                },
-                {"type": "text", "text": text_input},
-            ],
-        }
-    ]
-    
-    # Preparation for inference
-    text = processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-    image_inputs, video_inputs = process_vision_info(messages)
-    inputs = processor(
-        text=[text],
-        images=image_inputs,
-        videos=video_inputs,
-        padding=True,
-        return_tensors="pt",
-    )
-    inputs = inputs.to("cuda")
-    
-    # Inference: Generation of the output
-    generated_ids = model.generate(**inputs, max_new_tokens=1024)
-    generated_ids_trimmed = [
-        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-    ]
-    output_text = processor.batch_decode(
-        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-    )
-    
-    ocr_text = output_text[0]
-    
-    # If NER is enabled, process the OCR text
-    if run_ner:
-        ner_results = gliner_model.predict_entities(
-            ocr_text,
-            ner_labels.split(","),
-            threshold=0.3
+    try:
+        # First get the OCR text
+        text_input = "Convert the image to text."
+        image_path = array_to_image_path(image)
+        
+        model = models[model_id]
+        processor = processors[model_id]
+        
+        prompt = f"{user_prompt}<|image_1|>\n{text_input}{prompt_suffix}{assistant_prompt}"
+        image = Image.fromarray(image).convert("RGB")
+        messages = [
+        {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                        "image": image_path,
+                    },
+                    {"type": "text", "text": text_input},
+                ],
+            }
+        ]
+        
+        # Preparation for inference
+        text = processor.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        image_inputs, video_inputs = process_vision_info(messages)
+        inputs = processor(
+            text=[text],
+            images=image_inputs,
+            videos=video_inputs,
+            padding=True,
+            return_tensors="pt",
+        )
+        inputs = inputs.to("cuda")
+        
+        # Inference: Generation of the output
+        generated_ids = model.generate(**inputs, max_new_tokens=1024)
+        generated_ids_trimmed = [
+            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        output_text = processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
         )
         
-        # Create a list of tuples (text, label) for highlighting
-        highlighted_text = []
-        last_end = 0
+        ocr_text = output_text[0]
         
-        # Sort entities by start position
-        sorted_entities = sorted(ner_results, key=lambda x: x["start"])
+        # Create state dictionary
+        state_dict = {
+            "original_text": ocr_text,
+            "entities": []
+        }
         
-        # Process each entity and add non-entity text segments
-        for entity in sorted_entities:
-            # Add non-entity text before the current entity
-            if last_end < entity["start"]:
-                highlighted_text.append((ocr_text[last_end:entity["start"]], None))
+        # If NER is enabled, process the OCR text
+        if run_ner:
+            ner_results = gliner_model.predict_entities(
+                ocr_text,
+                ner_labels.split(","),
+                threshold=0.3
+            )
+            
+            # Update state with entities
+            state_dict["entities"] = ner_results
+            
+            # Create a list of tuples (text, label) for highlighting
+            highlighted_text = []
+            last_end = 0
+            
+            # Sort entities by start position
+            sorted_entities = sorted(ner_results, key=lambda x: x["start"])
+            
+            # Process each entity and add non-entity text segments
+            for entity in sorted_entities:
+                # Add non-entity text before the current entity
+                if last_end < entity["start"]:
+                    highlighted_text.append((ocr_text[last_end:entity["start"]], None))
+                
+                # Add the entity text with its label
+                highlighted_text.append((
+                    ocr_text[entity["start"]:entity["end"]],
+                    entity["label"]
+                ))
+                last_end = entity["end"]
+            
+            # Add any remaining text after the last entity
+            if last_end < len(ocr_text):
+                highlighted_text.append((ocr_text[last_end:], None))
             
-            # Add the entity text with its label
-            highlighted_text.append((
-                ocr_text[entity["start"]:entity["end"]],
-                entity["label"]
-            ))
-            last_end = entity["end"]
+            return highlighted_text, state_dict
         
-        # Add any remaining text after the last entity
-        if last_end < len(ocr_text):
-            highlighted_text.append((ocr_text[last_end:], None))
+        # If NER is disabled, return the text without highlighting
+        highlighted_text = [(ocr_text, None)]
+        return highlighted_text, state_dict
         
-        # Create TextWithMetadata instance with the highlighted text and metadata
-        result = TextWithMetadata(highlighted_text, original_text=ocr_text, entities=ner_results)
-        return result, result  # Return twice: once for display, once for state
-    
-    # If NER is disabled, return the text without highlighting
-    result = TextWithMetadata([(ocr_text, None)], original_text=ocr_text, entities=[])
-    return result, result  # Return twice: once for display, once for state
+    except Exception as e:
+        error_msg = f"Error processing image: {str(e)}"
+        return [(error_msg, None)], {"original_text": error_msg, "entities": []}
 
 css = """
   /* Overall app styling */
@@ -175,7 +181,6 @@ css = """
     padding: 20px;
     background-color: #f8f9fa;
   }
-
   /* Tabs styling */
   .tabs {
     border-radius: 8px;
@@ -183,7 +188,6 @@ css = """
     padding: 20px;
     box-shadow: 0 2px 6px rgba(0, 0, 0, 0.1);
   }
-
   /* Input/Output containers */
   .input-container, .output-container {
     background: white;
@@ -192,7 +196,6 @@ css = """
     margin: 10px 0;
     box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
   }
-
   /* Button styling */
   .submit-btn {
     background-color: #2d31fa !important;
@@ -207,7 +210,6 @@ css = """
     background-color: #1f24c7 !important;
     transform: translateY(-1px);
   }
-
   /* Output text area */
   #output {
     height: 500px;
@@ -218,13 +220,11 @@ css = """
     background: #ffffff;
     font-family: 'Arial', sans-serif;
   }
-
   /* Dropdown styling */
   .gr-dropdown {
     border-radius: 6px !important;
     border: 1px solid #e0e0e0 !important;
   }
-
   /* Image upload area */
   .gr-image-input {
     border: 2px dashed #ccc;
@@ -232,7 +232,6 @@ css = """
     padding: 20px;
     transition: all 0.3s ease;
   }
-
   .gr-image-input:hover {
     border-color: #2d31fa;
   }
@@ -283,7 +282,7 @@ with gr.Blocks(css=css) as demo:
     # Modify create_zip to use the state data
     def create_zip(image, fname, ocr_result):
         # Validate inputs
-        if not fname or image is None:  # Changed the validation check
+        if not fname or image is None:
             return None
         
         try:
@@ -298,9 +297,9 @@ with gr.Blocks(css=css) as demo:
                 img_path = os.path.join(temp_dir, f"{fname}.png")
                 image.save(img_path)
                 
-                # Use the OCR result from state
-                original_text = ocr_result.original_text if ocr_result else ""
-                entities = ocr_result.entities if ocr_result else []
+                # Use the OCR result from state - now it's a dictionary
+                original_text = ocr_result.get("original_text", "") if ocr_result else ""
+                entities = ocr_result.get("entities", []) if ocr_result else []
                 
                 # Save text
                 txt_path = os.path.join(temp_dir, f"{fname}.txt")