download ignores image

app.py

@@ -256,56 +256,64 @@ with gr.Blocks(css=css) as demo:
         download_output = gr.File(label="Download")
 
     def create_zip(image, text_data, fname):
-        if not isinstance(image, (Image.Image, np.ndarray)) or not text_data or not fname:
+        # Validate inputs
+        if not fname:
+            return None
+        if not text_data:
+            return None
+        if not isinstance(image, (Image.Image, np.ndarray)):
             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
-            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))
+        try:
+            # Convert numpy array to PIL Image if needed
+            if isinstance(image, np.ndarray):
+                image = Image.fromarray(image)
             
-            return zip_path
+            # Create a temporary directory
+            with tempfile.TemporaryDirectory() as temp_dir:
+                # Save image with default png extension
+                img_path = os.path.join(temp_dir, f"{fname}.png")
+                image.save(img_path)
+                
+                # Extract text content
+                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
+                
+                # 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 and entities
+                json_data = {
+                    "text": text,
+                    "entities": entities,
+                    "image_file": f"{fname}.png"
+                }
+                
+                # 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))
+                
+                # Read the zip file into memory
+                with open(zip_path, 'rb') as f:
+                    return (f"{fname}.zip", f.read())
+        except Exception as e:
+            print(f"Error creating zip: {str(e)}")
+            return None
 
     # Add the click handler for the download button
     download_btn.click(