fixed issues with image res

app.py

@@ -5,51 +5,59 @@ 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)
 
-
+# Prediction function
 def predict_from_url(url):
     prompt = "<OCR>"
     if not url:
-        return "Error: Please input a URL"
+        return "Error: Please input a URL", None
 
     try:
-        image = Image.open(BytesIO(requests.get(url).content))
+        # Open the image and convert to RGB format to handle grayscale or other formats
+        image = Image.open(BytesIO(requests.get(url).content)).convert("RGB")
     except Exception as e:
-        return f"Error: Failed to load image: {str(e)}"
+        return f"Error: Failed to load or process the image: {str(e)}", None
     
-    inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
-    generated_ids = model.generate(
-        input_ids=inputs["input_ids"],
-        pixel_values=inputs["pixel_values"],
-        max_new_tokens=4096,
-        num_beams=3,
-        do_sample=False
-    )
-    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-    parsed_answer = processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))
-
-
-    ocr_text = parsed_answer.get("<OCR>", "")
-
-    text = ocr_text.replace("\n", " ").strip()
-    output_text = " ".join(text.split())
+    # Preprocess and perform inference
+    try:
+        inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
+        generated_ids = model.generate(
+            input_ids=inputs["input_ids"],
+            pixel_values=inputs["pixel_values"],
+            max_new_tokens=4096,
+            num_beams=3,
+            do_sample=False
+        )
+        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+        parsed_answer = processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))
+        
+        # Extract OCR text
+        ocr_text = parsed_answer.get("<OCR>", "")
+        output_text = " ".join(ocr_text.replace("\n", " ").strip().split())
+    except Exception as e:
+        return f"Error: Failed to process the image for OCR: {str(e)}", image
 
     return output_text, image
 
+# Gradio Interface
 demo = gr.Interface(
-    fn=predict_from_url, 
+    fn=predict_from_url,
     inputs=gr.Textbox(label="Enter Image URL"),
     outputs=[
-        gr.Textbox(label="Output Text"),
-        gr.Image(label="Image")
+        gr.Textbox(label="Extracted Text"),
+        gr.Image(label="Uploaded Image")
     ],
-    title="OCR Text Extractor",
+    title="Enhanced OCR Text Extractor",
+    description="Provide an image URL, and this tool will extract text using OCR.",
     allow_flagging="never"
 )
 
+# Launch the app
 demo.launch()