Update app.py

app.py

@@ -1,33 +1,33 @@
 import torch
-import pytesseract
 import cv2
 import json
 import xml.etree.ElementTree as ET
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
-from layoutparser import Detectron2LayoutModel
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
-from stable_baselines3 import PPO
+import gradio as gr
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel, GPT2LMHeadModel, GPT2Tokenizer
 
-# Load OCR model
+# Load OCR model (TrOCR)
 processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
 model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
 
+# Load GPT-2 model
+GPT2_model = GPT2LMHeadModel.from_pretrained("gpt2")
+GPT2_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+
+# Image preprocessing
 def preprocess_image(image_path):
     image = cv2.imread(image_path)
     gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
     return gray
 
+# Extract text using TrOCR (instead of Tesseract)
 def extract_text(image_path):
     image = preprocess_image(image_path)
-    text = pytesseract.image_to_string(image)
+    pixel_values = processor(image, return_tensors="pt").pixel_values
+    generated_ids = model.generate(pixel_values)
+    text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
     return text
 
-def analyze_layout(image_path):
-    model = Detectron2LayoutModel("lp://PubLayNet/mask_rcnn_X_101_32x8d_FPN_3x/config")
-    image = cv2.imread(image_path)
-    layout = model.detect(image)
-    return layout
-
+# Generate structured format (JSON/XML)
 def generate_machine_readable_format(text, format_type='json'):
     if format_type == 'json':
         return json.dumps({"content": text})
@@ -38,41 +38,26 @@ def generate_machine_readable_format(text, format_type='json'):
         return ET.tostring(root, encoding='unicode')
     return text
 
-# Generative AI Model
-GPT2_model = GPT2LMHeadModel.from_pretrained("gpt2")
-GPT2_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-
+# GPT-2 for structured output
 def generate_structured_output(text):
     inputs = GPT2_tokenizer.encode(text, return_tensors="pt")
     outputs = GPT2_model.generate(inputs, max_length=500)
     return GPT2_tokenizer.decode(outputs[0])
 
-# Reinforcement Learning for Optimization
-class DocumentConversionEnv:
-    def __init__(self):
-        self.state = None
-    
-    def reset(self):
-        self.state = "start"
-        return self.state
-    
-    def step(self, action):
-        reward = 1 if action == "optimize" else -1
-        self.state = "optimized" if action == "optimize" else "start"
-        return self.state, reward, False, {}
-
-env = DocumentConversionEnv()
-rl_model = PPO("MlpPolicy", env, verbose=1)
-rl_model.learn(total_timesteps=1000)
-
-def convert_document(image_path, output_format='json'):
-    text = extract_text(image_path)
-    layout = analyze_layout(image_path)
+# Convert document
+def convert_document(image, output_format='json'):
+    text = extract_text(image)
     structured_output = generate_structured_output(text)
     machine_readable_output = generate_machine_readable_format(structured_output, format_type=output_format)
     return machine_readable_output
 
-# Example usage
-document_path = "sample_document.png"
-converted_document = convert_document(document_path, output_format='json')
-print(converted_document)
+# Gradio UI
+iface = gr.Interface(
+    fn=convert_document,
+    inputs=[gr.Image(type="filepath"), gr.Radio(["json", "xml"], label="Output Format")],
+    outputs="text",
+    title="Document OCR and Conversion",
+    description="Extracts text from images and converts it into structured JSON/XML format."
+)
+
+iface.launch()