app.py

import gradio as gr
import cv2
import numpy as np
import aiohttp
import asyncio
from stripeRemover import StripeRemover

class OCRUI:
    def __init__(self):
        self.API_URL = "http://s15.serv00.com:9081/compareAnalyze"
        self.stripe_remover = StripeRemover()

    def process_image(self, image, method):
        if image is None:
            return None
            
        try:
            if method == "Original":
                return cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
            elif method == "Fourier":
                return self.stripe_remover.fourier_method(image)
            elif method == "Morphological":
                return self.stripe_remover.morphological_method(image)
            elif method == "Adaptive":
                return self.stripe_remover.adaptive_threshold_method(image)
            elif method == "Enhanced":
                gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
                return self.stripe_remover.enhance_image(gray)
        except Exception as e:
            return None

    async def send_to_api(self, image_bytes):
        async with aiohttp.ClientSession() as session:
            data = aiohttp.FormData()
            data.add_field('image', image_bytes, 
                          filename='image.jpg',
                          content_type='image/jpeg')
            data.add_field('model', 'GEMINI')
            
            try:
                async with session.post(self.API_URL, data=data) as response:
                    return await response.json()
            except Exception as e:
                return {"error": str(e)}

    def ocr_process(self, image, task, preprocess_method):
        if image is None:
            return "Please upload an image", "No image provided"

        # Preprocess image
        processed_img = self.process_image(image, preprocess_method)
        if processed_img is None:
            return "Image processing failed", "Processing error"

        # Prepare image for API
        encode_params = [cv2.IMWRITE_JPEG_QUALITY, 50]
        _, img_bytes = cv2.imencode('.jpg', processed_img, encode_params)
        
        # Call API
        result = asyncio.run(self.send_to_api(img_bytes.tobytes()))
        
        if result is None or "error" in result:
            return "API call failed", "Error calling OCR service"
            
        # Format results
        text_output = f"Task: {task}\nResults:\n{str(result)}"
        html_output = f"<pre>{text_output}</pre>"
        
        return text_output, html_output

def create_ui():
    ui = OCRUI()
    
    with gr.Blocks() as demo:
        gr.Markdown("# 美宜佳DEMO")
        
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(type="numpy", label="Input Image")
                preprocess_dropdown = gr.Dropdown(
                    choices=["Original", "Fourier", "Morphological", "Adaptive", "Enhanced"],
                    label="Preprocessing Method",
                    value="Original"
                )
                task_dropdown = gr.Dropdown(
                    choices=["Plain OCR", "Format OCR", "Box OCR"],
                    label="OCR Task",
                    value="Plain OCR"
                )
                process_btn = gr.Button("Process Image")
                
            with gr.Column():
                text_output = gr.Textbox(label="OCR Results")
                html_output = gr.HTML(label="Formatted Results")

        process_btn.click(
            fn=ui.ocr_process,
            inputs=[image_input, task_dropdown, preprocess_dropdown],
            outputs=[text_output, html_output]
        )
    
    return demo

if __name__ == "__main__":
    demo = create_ui()
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)