Update app.py

app.py

@@ -1,107 +1,46 @@
-from typing import Optional
-import spaces
-
-import gradio as gr
-import numpy as np
-import torch
+from fastapi import FastAPI, UploadFile, Form
+from fastapi.responses import JSONResponse
+from transformers import AutoProcessor, AutoModelForVisualQuestionAnswering
 from PIL import Image
-import io
-
-
-import base64, os
-from utils import check_ocr_box, get_yolo_model, get_caption_model_processor, get_som_labeled_img
 import torch
-from PIL import Image
-
-# yolo_model = get_yolo_model(model_path='weights/icon_detect/best.pt')
-# caption_model_processor = get_caption_model_processor(model_name="florence2", model_name_or_path="weights/icon_caption_florence")
-
-from ultralytics import YOLO
-yolo_model = YOLO('best.pt').to('cpu')
-from transformers import AutoProcessor, AutoModelForCausalLM 
-processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("weights/icon_caption_florence", torch_dtype=torch.float16, trust_remote_code=True).to('cuda')
-caption_model_processor = {'processor': processor, 'model': model}
-print('finish loading model!!!')
-
-
-MARKDOWN = """
-# OmniParser for Pure Vision Based General GUI Agent 🔥
-<div>
-    <a href="https://arxiv.org/pdf/2408.00203">
-        <img src="https://img.shields.io/badge/arXiv-2408.00203-b31b1b.svg" alt="Arxiv" style="display:inline-block;">
-    </a>
-</div>
-
-OmniParser is a screen parsing tool to convert general GUI screen to structured elements. 
-
-📢 [[Project Page](https://microsoft.github.io/OmniParser/)] [[Blog Post](https://www.microsoft.com/en-us/research/articles/omniparser-for-pure-vision-based-gui-agent/)] [[Models](https://huggingface.co/microsoft/OmniParser)]
-"""
-
-# DEVICE = torch.device('cuda')
-
-@spaces.GPU
-@torch.inference_mode()
-# @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
-# @spaces.GPU(duration=65)
-def process(
-    image_input,
-    box_threshold,
-    iou_threshold
-) -> Optional[Image.Image]:
-
-    image_save_path = 'imgs/saved_image_demo.png'
-    image_input.save(image_save_path)
-    # import pdb; pdb.set_trace()
-    image = Image.open(image_save_path)
-    box_overlay_ratio = image.size[0] / 3200
-    draw_bbox_config = {
-        'text_scale': 0.8 * box_overlay_ratio,
-        'text_thickness': max(int(2 * box_overlay_ratio), 1),
-        'text_padding': max(int(3 * box_overlay_ratio), 1),
-        'thickness': max(int(3 * box_overlay_ratio), 1),
-    }
-
-    ocr_bbox_rslt, is_goal_filtered = check_ocr_box(image_save_path, display_img = False, output_bb_format='xyxy', goal_filtering=None, easyocr_args={'paragraph': False, 'text_threshold':0.9}, use_paddleocr=True)
-    text, ocr_bbox = ocr_bbox_rslt
-    # print('prompt:', prompt)
-    dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(image_save_path, yolo_model, BOX_TRESHOLD = box_threshold, output_coord_in_ratio=True, ocr_bbox=ocr_bbox,draw_bbox_config=draw_bbox_config, caption_model_processor=caption_model_processor, ocr_text=text,iou_threshold=iou_threshold)
-    image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))
-    print('finish processing')
-    parsed_content_list = '\n'.join(parsed_content_list)
-    return image, str(parsed_content_list), str(label_coordinates)
-
-
-
-with gr.Blocks() as demo:
-    gr.Markdown(MARKDOWN)
-    with gr.Row():
-        with gr.Column():
-            image_input_component = gr.Image(
-                type='pil', label='Upload image')
-            # set the threshold for removing the bounding boxes with low confidence, default is 0.05
-            box_threshold_component = gr.Slider(
-                label='Box Threshold', minimum=0.01, maximum=1.0, step=0.01, value=0.05)
-            # set the threshold for removing the bounding boxes with large overlap, default is 0.1
-            iou_threshold_component = gr.Slider(
-                label='IOU Threshold', minimum=0.01, maximum=1.0, step=0.01, value=0.1)
-            submit_button_component = gr.Button(
-                value='Submit', variant='primary')
-        with gr.Column():
-            image_output_component = gr.Image(type='pil', label='Image Output')
-            text_output_component = gr.Textbox(label='Parsed screen elements', placeholder='Text Output')
-            coordinates_output_component = gr.Textbox(label='Coordinates', placeholder='Coordinates Output')
-
-    submit_button_component.click(
-        fn=process,
-        inputs=[
-            image_input_component,
-            box_threshold_component,
-            iou_threshold_component
-        ],
-        outputs=[image_output_component, text_output_component, coordinates_output_component]
-    )
-
-# demo.launch(debug=False, show_error=True, share=True)
-# demo.launch(share=True, server_port=7861, server_name='0.0.0.0')
-demo.queue().launch(share=False)
+import uvicorn
+
+# Initialize FastAPI app
+app = FastAPI()
+
+# Load model and processor with trust_remote_code=True
+processor = AutoProcessor.from_pretrained("Sanket17/hello", trust_remote_code=True)
+model = AutoModelForVisualQuestionAnswering.from_pretrained("Sanket17/hello", trust_remote_code=True)
+
+@app.post("/vqa/")
+async def visual_question_answer(file: UploadFile, question: str = Form(...)):
+    """
+    Endpoint for visual question answering.
+    - file: Upload an image file
+    - question: Textual question about the image
+    """
+    try:
+        # Load image
+        image = Image.open(file.file).convert("RGB")
+        
+        # Preprocess inputs
+        inputs = processor(images=image, text=question, return_tensors="pt")
+        
+        # Get model predictions
+        outputs = model(**inputs)
+        
+        # Decode the answer (check model output for correct handling)
+        answer = outputs.logits.argmax(dim=-1).item()  # Example way to get the answer index
+        
+        # If the output logits contain a mapping, we can return the answer string
+        answer_str = processor.decode([answer])  # Assuming you get the answer index from logits
+
+        # Return JSON response
+        return JSONResponse(content={"question": question, "answer": answer_str})
+    
+    except Exception as e:
+        return JSONResponse(content={"error": str(e)}, status_code=500)
+
+# Start the FastAPI server
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)