Update app.py

app.py

@@ -34,7 +34,12 @@ def run_example(task_prompt, image):
 
 def process_image(image):
     image = Image.fromarray(image)  # Convert NumPy array to PIL Image
-    task_prompt = '<PC>'
+    if task_prompt == 'Product Caption':
+        task_prompt = '<PC>'
+        results = run_example(task_prompt, image, model_id=model_id)
+    elif task_prompt == 'OCR':
+        task_prompt = '<OCR>'
+        results = run_example(task_prompt, image, model_id=model_id)
     results = run_example(task_prompt, image)
     
     # Remove the key and get the text value
@@ -71,13 +76,16 @@ document.querySelector('button').addEventListener('click', function() {
     setTimeout(adjustHeight, 500);  // Adjust the height after a small delay to ensure content is loaded
 });
 """
-
+single_task_list =[
+    'Product Caption', 'OCR'
+]
 with gr.Blocks(css=css) as demo:
     gr.Markdown(DESCRIPTION)
     with gr.Tab(label="Product Image Select"):
         with gr.Row():
             with gr.Column():
                 input_img = gr.Image(label="Input Picture")
+                task_prompt = gr.Dropdown(choices=single_task_list, label="Task Prompt", value="Product Caption")
                 submit_btn = gr.Button(value="Submit")
             with gr.Column():
                 output_text = gr.HTML(label="Output Text", elem_id="output")