update app.py

app.py

@@ -6,22 +6,23 @@ import base64
 import os
 from dotenv import load_dotenv
 from openai import OpenAI
+from PIL import Image
+import io
 
 # Load environment variables
 load_dotenv()
 
 # Function to upload image to imgbb
-def upload_image_to_imgbb(image_path):
+def upload_image_to_imgbb(image_bytes):
     try:
-        with open(image_path, "rb") as file:
-            url = "https://api.imgbb.com/1/upload"
-            payload = {
-                "key": os.getenv("IMGBB_API_KEY"),
-                "image": base64.b64encode(file.read()).decode()
-            }
-            res = requests.post(url, data=payload)
-            res.raise_for_status()  # Raises an HTTPError for bad responses
-            return res.json()['data']['url']
+        url = "https://api.imgbb.com/1/upload"
+        payload = {
+            "key": os.getenv("IMGBB_API_KEY"),
+            "image": base64.b64encode(image_bytes).decode()
+        }
+        res = requests.post(url, data=payload)
+        res.raise_for_status()  # Raises an HTTPError for bad responses
+        return res.json()['data']['url']
     except Exception as e:
         return f"Error uploading image: {str(e)}"
 
@@ -54,7 +55,7 @@ def process_ocr_with_gpt(ocr_results):
 
         {ocr_results}
 
-        Please process this information and provide a clean, well-formatted output. If there are any obvious errors in the OCR, please correct them. Format the output in a clear and readable manner.
+        Please process this information and provide a clean, well-formatted output. Arrange all elements in order and omit any elements not present in the file.
         """
 
         response = openai_client.chat.completions.create(
@@ -70,42 +71,64 @@ def process_ocr_with_gpt(ocr_results):
         return f"Error in GPT processing: {str(e)}"
 
 # Gradio interface function
-def process_image(image):
-    try:
-        # Save the image temporarily and get its path
-        image_path = "temp_image.png"
-        image.save(image_path)
-        
-        # Upload image to imgbb and get URL
-        image_url = upload_image_to_imgbb(image_path)
-        if image_url.startswith("Error"):
-            return image_url, "Failed to process due to image upload error"
-        
-        # Get OCR results
-        ocr_results = get_ocr_results(image_url)
-        if ocr_results.startswith("Error"):
-            return ocr_results, "Failed to process due to OCR error"
-        
-        # Process with GPT
-        processed_results = process_ocr_with_gpt(ocr_results)
-        
-        # Clean up temporary file
-        os.remove(image_path)
-        
-        return ocr_results, processed_results
-    except Exception as e:
-        return f"Error in image processing: {str(e)}", "Failed to process due to an error"
+def process_images(files):
+    ocr_results_list = []
+    processed_results_list = []
+    download_links = []
+    
+    for idx, file in enumerate(files):
+        try:
+            # Read the uploaded binary file
+            image = Image.open(io.BytesIO(file))
+            image_path = f"temp_image_{idx}.png"
+            image.save(image_path)
+            
+            # Upload image to imgbb and get URL
+            image_url = upload_image_to_imgbb(file)
+            if image_url.startswith("Error"):
+                ocr_results_list.append(image_url)
+                processed_results_list.append("Failed to process due to image upload error")
+                continue
+            
+            # Get OCR results
+            ocr_results = get_ocr_results(image_url)
+            if ocr_results.startswith("Error"):
+                ocr_results_list.append(ocr_results)
+                processed_results_list.append("Failed to process due to OCR error")
+                continue
+            
+            # Process with GPT
+            processed_results = process_ocr_with_gpt(ocr_results)
+            
+            # Save processed results to a file for download
+            result_file_path = f"processed_result_{idx}.txt"
+            with open(result_file_path, 'w') as result_file:
+                result_file.write(processed_results)
+            
+            download_links.append(result_file_path)
+            
+            # Clean up temporary file
+            os.remove(image_path)
+            
+            ocr_results_list.append(ocr_results)
+            processed_results_list.append(processed_results)
+        except Exception as e:
+            ocr_results_list.append(f"Error in image processing: {str(e)}")
+            processed_results_list.append("Failed to process due to an error")
+
+    return ocr_results_list, processed_results_list, download_links
 
 # Create Gradio interface
 iface = gr.Interface(
-    fn=process_image,
-    inputs=gr.Image(type="pil"),
+    fn=process_images,
+    inputs=gr.Files(label="Upload Images", file_count="multiple", type="binary"),
     outputs=[
         gr.Textbox(label="OCR Results"),
-        gr.Textbox(label="Processed Results")
+        gr.Textbox(label="Processed Results"),
+        gr.File(label="Download Processed Results")
     ],
     title="OCR and Text Processing App",
-    description="Upload an image to extract text and process it."
+    description="Upload images to extract text and process it. Download the processed results."
 )
 
 # Launch the app