updated ocr

app.py

@@ -1,31 +1,26 @@
 import gradio as gr
 import requests
-from PIL import Image, ImageDraw
+from PIL import Image
 from transformers import AutoProcessor, AutoModelForCausalLM
 from io import BytesIO
 import torch
 
-# Set device
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# Load model and processor
 model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-large", torch_dtype=torch_dtype, trust_remote_code=True).to(device)
 processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True)
 
-# List of colors to cycle through for bounding boxes
-COLORS = ["red", "blue", "green", "yellow", "purple", "orange", "cyan", "magenta"]
 
-# Prediction function
 def predict_from_url(url):
-    prompt = "<OD>"
+    prompt = "<OCR>"
     if not url:
-        return {"Error": "Please input a URL"}, None
+        return "Error: Please input a URL"
 
     try:
         image = Image.open(BytesIO(requests.get(url).content))
     except Exception as e:
-        return {"Error": f"Failed to load image: {str(e)}"}, None
+        return f"Error: Failed to load image: {str(e)}"
     
     inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
     generated_ids = model.generate(
@@ -36,34 +31,25 @@ def predict_from_url(url):
         do_sample=False
     )
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-    parsed_answer = processor.post_process_generation(generated_text, task="<OD>", image_size=(image.width, image.height))
+    parsed_answer = processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))
 
-    labels = parsed_answer.get('<OD>', {}).get('labels', [])
-    bboxes = parsed_answer.get('<OD>', {}).get('bboxes', [])
-    
-    # Draw bounding boxes on the image
-    draw = ImageDraw.Draw(image)
-    legend = []  # Store legend entries
-    for idx, (bbox, label) in enumerate(zip(bboxes, labels)):
-        x1, y1, x2, y2 = bbox
-        color = COLORS[idx % len(COLORS)]  # Cycle through colors
-        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
-        legend.append(f"{label}: {color}")
-    
 
-    return "\n".join(legend), image
+    ocr_text = parsed_answer.get("<OCR>", "")
+
+    text = ocr_text.replace("\n", " ").strip()
+    output_text = " ".join(text.split())
+
+    return output_text, image
 
-# Gradio interface
 demo = gr.Interface(
     fn=predict_from_url, 
     inputs=gr.Textbox(label="Enter Image URL"),
     outputs=[
-        gr.Textbox(label="Legend"),  # Output the legend
-        gr.Image(label="Image with Bounding Boxes")  # Output the processed image
+        gr.Textbox(label="Output Text"),
+        gr.Image(label="Image")
     ],
-    title="Item Classifier with Bounding Boxes and Legend",
+    title="OCR Text Extractor",
     allow_flagging="never"
 )
 
-# Launch the interface
 demo.launch()