Spaces:
Sleeping
Sleeping
[email protected]
commited on
Commit
·
3381ab5
0
Parent(s):
Added model files
Browse files- .DS_Store +0 -0
- .gitattributes +2 -0
- Dockerfile +32 -0
- README.md +10 -0
- app.py +132 -0
- download_models.py +50 -0
- requirements.txt +21 -0
- utils.py +416 -0
.DS_Store
ADDED
Binary file (6.15 kB). View file
|
|
.gitattributes
ADDED
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
1 |
+
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
2 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
Dockerfile
ADDED
@@ -0,0 +1,32 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
FROM python:3.9-slim
|
2 |
+
|
3 |
+
# Install system dependencies
|
4 |
+
RUN apt-get update && apt-get install -y \
|
5 |
+
libgl1-mesa-glx \
|
6 |
+
libglib2.0-0 \
|
7 |
+
wget \
|
8 |
+
git \
|
9 |
+
&& rm -rf /var/lib/apt/lists/*
|
10 |
+
|
11 |
+
WORKDIR /app
|
12 |
+
|
13 |
+
# Copy requirements and install dependencies
|
14 |
+
COPY requirements.txt .
|
15 |
+
RUN pip install --no-cache-dir -r requirements.txt
|
16 |
+
|
17 |
+
# Create directories for models
|
18 |
+
RUN mkdir -p weights/icon_detect weights/icon_caption_florence
|
19 |
+
|
20 |
+
# Copy application code
|
21 |
+
COPY main.py utils.py download_models.py ./
|
22 |
+
|
23 |
+
# Download models during build
|
24 |
+
RUN python download_models.py
|
25 |
+
|
26 |
+
# Set up user for HF Spaces
|
27 |
+
RUN useradd -m -u 1000 user
|
28 |
+
USER user
|
29 |
+
ENV PATH="/home/user/.local/bin:$PATH"
|
30 |
+
|
31 |
+
# Run the application
|
32 |
+
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
|
README.md
ADDED
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
---
|
2 |
+
title: Parser Own
|
3 |
+
emoji: 💻
|
4 |
+
colorFrom: blue
|
5 |
+
colorTo: gray
|
6 |
+
sdk: docker
|
7 |
+
pinned: false
|
8 |
+
---
|
9 |
+
|
10 |
+
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
app.py
ADDED
@@ -0,0 +1,132 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from fastapi import FastAPI, File, UploadFile, HTTPException
|
2 |
+
from fastapi.responses import JSONResponse
|
3 |
+
from pydantic import BaseModel
|
4 |
+
from typing import Optional
|
5 |
+
import base64
|
6 |
+
import io
|
7 |
+
from PIL import Image
|
8 |
+
import torch
|
9 |
+
import numpy as np
|
10 |
+
import os
|
11 |
+
|
12 |
+
# Existing imports
|
13 |
+
import numpy as np
|
14 |
+
import torch
|
15 |
+
from PIL import Image
|
16 |
+
import io
|
17 |
+
|
18 |
+
from utils import (
|
19 |
+
check_ocr_box,
|
20 |
+
get_yolo_model,
|
21 |
+
get_caption_model_processor,
|
22 |
+
get_som_labeled_img,
|
23 |
+
)
|
24 |
+
import torch
|
25 |
+
|
26 |
+
# yolo_model = get_yolo_model(model_path='/data/icon_detect/best.pt')
|
27 |
+
# caption_model_processor = get_caption_model_processor(model_name="florence2", model_name_or_path="/data/icon_caption_florence")
|
28 |
+
|
29 |
+
from ultralytics import YOLO
|
30 |
+
|
31 |
+
# if not os.path.exists("/data/icon_detect"):
|
32 |
+
# os.makedirs("/data/icon_detect")
|
33 |
+
|
34 |
+
try:
|
35 |
+
yolo_model = YOLO("weights/icon_detect/best.pt").to("cuda")
|
36 |
+
except:
|
37 |
+
yolo_model = YOLO("weights/icon_detect/best.pt")
|
38 |
+
|
39 |
+
from transformers import AutoProcessor, AutoModelForCausalLM
|
40 |
+
|
41 |
+
processor = AutoProcessor.from_pretrained(
|
42 |
+
"microsoft/Florence-2-base", trust_remote_code=True
|
43 |
+
)
|
44 |
+
|
45 |
+
try:
|
46 |
+
model = AutoModelForCausalLM.from_pretrained(
|
47 |
+
"weights/icon_caption_florence",
|
48 |
+
torch_dtype=torch.float16,
|
49 |
+
trust_remote_code=True,
|
50 |
+
).to("cuda")
|
51 |
+
except:
|
52 |
+
model = AutoModelForCausalLM.from_pretrained(
|
53 |
+
"weights/icon_caption_florence",
|
54 |
+
torch_dtype=torch.float16,
|
55 |
+
trust_remote_code=True,
|
56 |
+
)
|
57 |
+
caption_model_processor = {"processor": processor, "model": model}
|
58 |
+
print("finish loading model!!!")
|
59 |
+
|
60 |
+
app = FastAPI()
|
61 |
+
|
62 |
+
|
63 |
+
class ProcessResponse(BaseModel):
|
64 |
+
image: str # Base64 encoded image
|
65 |
+
parsed_content_list: str
|
66 |
+
label_coordinates: str
|
67 |
+
|
68 |
+
|
69 |
+
def process(
|
70 |
+
image_input: Image.Image, box_threshold: float, iou_threshold: float
|
71 |
+
) -> ProcessResponse:
|
72 |
+
image_save_path = "imgs/saved_image_demo.png"
|
73 |
+
image_input.save(image_save_path)
|
74 |
+
image = Image.open(image_save_path)
|
75 |
+
box_overlay_ratio = image.size[0] / 3200
|
76 |
+
draw_bbox_config = {
|
77 |
+
"text_scale": 0.8 * box_overlay_ratio,
|
78 |
+
"text_thickness": max(int(2 * box_overlay_ratio), 1),
|
79 |
+
"text_padding": max(int(3 * box_overlay_ratio), 1),
|
80 |
+
"thickness": max(int(3 * box_overlay_ratio), 1),
|
81 |
+
}
|
82 |
+
|
83 |
+
ocr_bbox_rslt, is_goal_filtered = check_ocr_box(
|
84 |
+
image_save_path,
|
85 |
+
display_img=False,
|
86 |
+
output_bb_format="xyxy",
|
87 |
+
goal_filtering=None,
|
88 |
+
easyocr_args={"paragraph": False, "text_threshold": 0.9},
|
89 |
+
use_paddleocr=True,
|
90 |
+
)
|
91 |
+
text, ocr_bbox = ocr_bbox_rslt
|
92 |
+
dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(
|
93 |
+
image_save_path,
|
94 |
+
yolo_model,
|
95 |
+
BOX_TRESHOLD=box_threshold,
|
96 |
+
output_coord_in_ratio=True,
|
97 |
+
ocr_bbox=ocr_bbox,
|
98 |
+
draw_bbox_config=draw_bbox_config,
|
99 |
+
caption_model_processor=caption_model_processor,
|
100 |
+
ocr_text=text,
|
101 |
+
iou_threshold=iou_threshold,
|
102 |
+
)
|
103 |
+
image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))
|
104 |
+
print("finish processing")
|
105 |
+
parsed_content_list_str = "\n".join(parsed_content_list)
|
106 |
+
|
107 |
+
# Encode image to base64
|
108 |
+
buffered = io.BytesIO()
|
109 |
+
image.save(buffered, format="PNG")
|
110 |
+
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
111 |
+
|
112 |
+
return ProcessResponse(
|
113 |
+
image=img_str,
|
114 |
+
parsed_content_list=str(parsed_content_list_str),
|
115 |
+
label_coordinates=str(label_coordinates),
|
116 |
+
)
|
117 |
+
|
118 |
+
|
119 |
+
@app.post("/process_image", response_model=ProcessResponse)
|
120 |
+
async def process_image(
|
121 |
+
image_file: UploadFile = File(...),
|
122 |
+
box_threshold: float = 0.05,
|
123 |
+
iou_threshold: float = 0.1,
|
124 |
+
):
|
125 |
+
try:
|
126 |
+
contents = await image_file.read()
|
127 |
+
image_input = Image.open(io.BytesIO(contents)).convert("RGB")
|
128 |
+
except Exception as e:
|
129 |
+
raise HTTPException(status_code=400, detail="Invalid image file")
|
130 |
+
|
131 |
+
response = process(image_input, box_threshold, iou_threshold)
|
132 |
+
return response
|
download_models.py
ADDED
@@ -0,0 +1,50 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
from huggingface_hub import hf_hub_download
|
3 |
+
|
4 |
+
def download_models():
|
5 |
+
# Create directories if they don't exist
|
6 |
+
os.makedirs("weights/icon_detect", exist_ok=True)
|
7 |
+
os.makedirs("weights/icon_caption_florence", exist_ok=True)
|
8 |
+
|
9 |
+
# Files to download from banao-tech/OmniParser
|
10 |
+
florence_files = [
|
11 |
+
"weights/icon_caption_florence/config.json",
|
12 |
+
"weights/icon_caption_florence/generation_config.json",
|
13 |
+
"weights/icon_caption_florence/model.safetensors"
|
14 |
+
]
|
15 |
+
|
16 |
+
yolo_files = [
|
17 |
+
"weights/icon_detect/best.pt",
|
18 |
+
"weights/icon_detect/model.yaml"
|
19 |
+
]
|
20 |
+
|
21 |
+
# Download Florence model files
|
22 |
+
for file_path in florence_files:
|
23 |
+
if not os.path.exists(file_path):
|
24 |
+
print(f"Downloading {file_path}...")
|
25 |
+
try:
|
26 |
+
hf_hub_download(
|
27 |
+
repo_id="banao-tech/OmniParser",
|
28 |
+
filename=file_path,
|
29 |
+
local_dir="."
|
30 |
+
)
|
31 |
+
print(f"Successfully downloaded {file_path}")
|
32 |
+
except Exception as e:
|
33 |
+
print(f"Error downloading {file_path}: {str(e)}")
|
34 |
+
|
35 |
+
# Download YOLO model files
|
36 |
+
for file_path in yolo_files:
|
37 |
+
if not os.path.exists(file_path):
|
38 |
+
print(f"Downloading {file_path}...")
|
39 |
+
try:
|
40 |
+
hf_hub_download(
|
41 |
+
repo_id="banao-tech/OmniParser",
|
42 |
+
filename=file_path,
|
43 |
+
local_dir="."
|
44 |
+
)
|
45 |
+
print(f"Successfully downloaded {file_path}")
|
46 |
+
except Exception as e:
|
47 |
+
print(f"Error downloading {file_path}: {str(e)}")
|
48 |
+
|
49 |
+
if __name__ == "__main__":
|
50 |
+
download_models()
|
requirements.txt
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
torch
|
2 |
+
easyocr
|
3 |
+
torchvision
|
4 |
+
supervision==0.18.0
|
5 |
+
openai==1.3.5
|
6 |
+
transformers
|
7 |
+
ultralytics==8.1.24
|
8 |
+
azure-identity
|
9 |
+
numpy
|
10 |
+
opencv-python
|
11 |
+
opencv-python-headless
|
12 |
+
gradio
|
13 |
+
dill
|
14 |
+
accelerate
|
15 |
+
timm
|
16 |
+
einops==0.8.0
|
17 |
+
paddlepaddle
|
18 |
+
paddleocr
|
19 |
+
fastapi
|
20 |
+
uvicorn
|
21 |
+
huggingface_hub
|
utils.py
ADDED
@@ -0,0 +1,416 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# from ultralytics import YOLO
|
2 |
+
import os
|
3 |
+
import io
|
4 |
+
import base64
|
5 |
+
import time
|
6 |
+
from PIL import Image, ImageDraw, ImageFont
|
7 |
+
import json
|
8 |
+
import requests
|
9 |
+
# utility function
|
10 |
+
import os
|
11 |
+
|
12 |
+
import json
|
13 |
+
import sys
|
14 |
+
import os
|
15 |
+
import cv2
|
16 |
+
import numpy as np
|
17 |
+
# %matplotlib inline
|
18 |
+
from matplotlib import pyplot as plt
|
19 |
+
import easyocr
|
20 |
+
from paddleocr import PaddleOCR
|
21 |
+
reader = easyocr.Reader(['en'])
|
22 |
+
paddle_ocr = PaddleOCR(
|
23 |
+
lang='en', # other lang also available
|
24 |
+
use_angle_cls=False,
|
25 |
+
use_gpu=False, # using cuda will conflict with pytorch in the same process
|
26 |
+
show_log=False,
|
27 |
+
max_batch_size=1024,
|
28 |
+
use_dilation=True, # improves accuracy
|
29 |
+
det_db_score_mode='slow', # improves accuracy
|
30 |
+
rec_batch_num=1024)
|
31 |
+
import time
|
32 |
+
import base64
|
33 |
+
|
34 |
+
import os
|
35 |
+
import ast
|
36 |
+
import torch
|
37 |
+
from typing import Tuple, List
|
38 |
+
from torchvision.ops import box_convert
|
39 |
+
import re
|
40 |
+
from torchvision.transforms import ToPILImage
|
41 |
+
import supervision as sv
|
42 |
+
import torchvision.transforms as T
|
43 |
+
|
44 |
+
|
45 |
+
def get_caption_model_processor(model_name, model_name_or_path="Salesforce/blip2-opt-2.7b", device=None):
|
46 |
+
if not device:
|
47 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
48 |
+
if model_name == "blip2":
|
49 |
+
from transformers import Blip2Processor, Blip2ForConditionalGeneration
|
50 |
+
processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
|
51 |
+
if device == 'cpu':
|
52 |
+
model = Blip2ForConditionalGeneration.from_pretrained(
|
53 |
+
model_name_or_path, device_map=None, torch_dtype=torch.float32
|
54 |
+
)
|
55 |
+
else:
|
56 |
+
model = Blip2ForConditionalGeneration.from_pretrained(
|
57 |
+
model_name_or_path, device_map=None, torch_dtype=torch.float16
|
58 |
+
).to(device)
|
59 |
+
elif model_name == "florence2":
|
60 |
+
from transformers import AutoProcessor, AutoModelForCausalLM
|
61 |
+
processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
|
62 |
+
if device == 'cpu':
|
63 |
+
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch.float32, trust_remote_code=True)
|
64 |
+
else:
|
65 |
+
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch.float16, trust_remote_code=True).to(device)
|
66 |
+
return {'model': model.to(device), 'processor': processor}
|
67 |
+
|
68 |
+
|
69 |
+
def get_yolo_model(model_path):
|
70 |
+
from ultralytics import YOLO
|
71 |
+
# Load the model.
|
72 |
+
model = YOLO(model_path)
|
73 |
+
return model
|
74 |
+
|
75 |
+
|
76 |
+
@torch.inference_mode()
|
77 |
+
def get_parsed_content_icon(filtered_boxes, ocr_bbox, image_source, caption_model_processor, prompt=None):
|
78 |
+
to_pil = ToPILImage()
|
79 |
+
if ocr_bbox:
|
80 |
+
non_ocr_boxes = filtered_boxes[len(ocr_bbox):]
|
81 |
+
else:
|
82 |
+
non_ocr_boxes = filtered_boxes
|
83 |
+
croped_pil_image = []
|
84 |
+
for i, coord in enumerate(non_ocr_boxes):
|
85 |
+
xmin, xmax = int(coord[0]*image_source.shape[1]), int(coord[2]*image_source.shape[1])
|
86 |
+
ymin, ymax = int(coord[1]*image_source.shape[0]), int(coord[3]*image_source.shape[0])
|
87 |
+
cropped_image = image_source[ymin:ymax, xmin:xmax, :]
|
88 |
+
croped_pil_image.append(to_pil(cropped_image))
|
89 |
+
|
90 |
+
model, processor = caption_model_processor['model'], caption_model_processor['processor']
|
91 |
+
if not prompt:
|
92 |
+
if 'florence' in model.config.name_or_path:
|
93 |
+
prompt = "<CAPTION>"
|
94 |
+
else:
|
95 |
+
prompt = "The image shows"
|
96 |
+
|
97 |
+
batch_size = 10 # Number of samples per batch
|
98 |
+
generated_texts = []
|
99 |
+
device = model.device
|
100 |
+
|
101 |
+
for i in range(0, len(croped_pil_image), batch_size):
|
102 |
+
batch = croped_pil_image[i:i+batch_size]
|
103 |
+
if model.device.type == 'cuda':
|
104 |
+
inputs = processor(images=batch, text=[prompt]*len(batch), return_tensors="pt").to(device=device, dtype=torch.float16)
|
105 |
+
else:
|
106 |
+
inputs = processor(images=batch, text=[prompt]*len(batch), return_tensors="pt").to(device=device)
|
107 |
+
if 'florence' in model.config.name_or_path:
|
108 |
+
generated_ids = model.generate(input_ids=inputs["input_ids"],pixel_values=inputs["pixel_values"],max_new_tokens=1024,num_beams=3, do_sample=False)
|
109 |
+
else:
|
110 |
+
generated_ids = model.generate(**inputs, max_length=100, num_beams=5, no_repeat_ngram_size=2, early_stopping=True, num_return_sequences=1) # temperature=0.01, do_sample=True,
|
111 |
+
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
|
112 |
+
generated_text = [gen.strip() for gen in generated_text]
|
113 |
+
generated_texts.extend(generated_text)
|
114 |
+
|
115 |
+
return generated_texts
|
116 |
+
|
117 |
+
|
118 |
+
|
119 |
+
def get_parsed_content_icon_phi3v(filtered_boxes, ocr_bbox, image_source, caption_model_processor):
|
120 |
+
to_pil = ToPILImage()
|
121 |
+
if ocr_bbox:
|
122 |
+
non_ocr_boxes = filtered_boxes[len(ocr_bbox):]
|
123 |
+
else:
|
124 |
+
non_ocr_boxes = filtered_boxes
|
125 |
+
croped_pil_image = []
|
126 |
+
for i, coord in enumerate(non_ocr_boxes):
|
127 |
+
xmin, xmax = int(coord[0]*image_source.shape[1]), int(coord[2]*image_source.shape[1])
|
128 |
+
ymin, ymax = int(coord[1]*image_source.shape[0]), int(coord[3]*image_source.shape[0])
|
129 |
+
cropped_image = image_source[ymin:ymax, xmin:xmax, :]
|
130 |
+
croped_pil_image.append(to_pil(cropped_image))
|
131 |
+
|
132 |
+
model, processor = caption_model_processor['model'], caption_model_processor['processor']
|
133 |
+
device = model.device
|
134 |
+
messages = [{"role": "user", "content": "<|image_1|>\ndescribe the icon in one sentence"}]
|
135 |
+
prompt = processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
136 |
+
|
137 |
+
batch_size = 5 # Number of samples per batch
|
138 |
+
generated_texts = []
|
139 |
+
|
140 |
+
for i in range(0, len(croped_pil_image), batch_size):
|
141 |
+
images = croped_pil_image[i:i+batch_size]
|
142 |
+
image_inputs = [processor.image_processor(x, return_tensors="pt") for x in images]
|
143 |
+
inputs ={'input_ids': [], 'attention_mask': [], 'pixel_values': [], 'image_sizes': []}
|
144 |
+
texts = [prompt] * len(images)
|
145 |
+
for i, txt in enumerate(texts):
|
146 |
+
input = processor._convert_images_texts_to_inputs(image_inputs[i], txt, return_tensors="pt")
|
147 |
+
inputs['input_ids'].append(input['input_ids'])
|
148 |
+
inputs['attention_mask'].append(input['attention_mask'])
|
149 |
+
inputs['pixel_values'].append(input['pixel_values'])
|
150 |
+
inputs['image_sizes'].append(input['image_sizes'])
|
151 |
+
max_len = max([x.shape[1] for x in inputs['input_ids']])
|
152 |
+
for i, v in enumerate(inputs['input_ids']):
|
153 |
+
inputs['input_ids'][i] = torch.cat([processor.tokenizer.pad_token_id * torch.ones(1, max_len - v.shape[1], dtype=torch.long), v], dim=1)
|
154 |
+
inputs['attention_mask'][i] = torch.cat([torch.zeros(1, max_len - v.shape[1], dtype=torch.long), inputs['attention_mask'][i]], dim=1)
|
155 |
+
inputs_cat = {k: torch.concatenate(v).to(device) for k, v in inputs.items()}
|
156 |
+
|
157 |
+
generation_args = {
|
158 |
+
"max_new_tokens": 25,
|
159 |
+
"temperature": 0.01,
|
160 |
+
"do_sample": False,
|
161 |
+
}
|
162 |
+
generate_ids = model.generate(**inputs_cat, eos_token_id=processor.tokenizer.eos_token_id, **generation_args)
|
163 |
+
# # remove input tokens
|
164 |
+
generate_ids = generate_ids[:, inputs_cat['input_ids'].shape[1]:]
|
165 |
+
response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
|
166 |
+
response = [res.strip('\n').strip() for res in response]
|
167 |
+
generated_texts.extend(response)
|
168 |
+
|
169 |
+
return generated_texts
|
170 |
+
|
171 |
+
def remove_overlap(boxes, iou_threshold, ocr_bbox=None):
|
172 |
+
assert ocr_bbox is None or isinstance(ocr_bbox, List)
|
173 |
+
|
174 |
+
def box_area(box):
|
175 |
+
return (box[2] - box[0]) * (box[3] - box[1])
|
176 |
+
|
177 |
+
def intersection_area(box1, box2):
|
178 |
+
x1 = max(box1[0], box2[0])
|
179 |
+
y1 = max(box1[1], box2[1])
|
180 |
+
x2 = min(box1[2], box2[2])
|
181 |
+
y2 = min(box1[3], box2[3])
|
182 |
+
return max(0, x2 - x1) * max(0, y2 - y1)
|
183 |
+
|
184 |
+
def IoU(box1, box2):
|
185 |
+
intersection = intersection_area(box1, box2)
|
186 |
+
union = box_area(box1) + box_area(box2) - intersection + 1e-6
|
187 |
+
if box_area(box1) > 0 and box_area(box2) > 0:
|
188 |
+
ratio1 = intersection / box_area(box1)
|
189 |
+
ratio2 = intersection / box_area(box2)
|
190 |
+
else:
|
191 |
+
ratio1, ratio2 = 0, 0
|
192 |
+
return max(intersection / union, ratio1, ratio2)
|
193 |
+
|
194 |
+
boxes = boxes.tolist()
|
195 |
+
filtered_boxes = []
|
196 |
+
if ocr_bbox:
|
197 |
+
filtered_boxes.extend(ocr_bbox)
|
198 |
+
# print('ocr_bbox!!!', ocr_bbox)
|
199 |
+
for i, box1 in enumerate(boxes):
|
200 |
+
# if not any(IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2) for j, box2 in enumerate(boxes) if i != j):
|
201 |
+
is_valid_box = True
|
202 |
+
for j, box2 in enumerate(boxes):
|
203 |
+
if i != j and IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2):
|
204 |
+
is_valid_box = False
|
205 |
+
break
|
206 |
+
if is_valid_box:
|
207 |
+
# add the following 2 lines to include ocr bbox
|
208 |
+
if ocr_bbox:
|
209 |
+
if not any(IoU(box1, box3) > iou_threshold for k, box3 in enumerate(ocr_bbox)):
|
210 |
+
filtered_boxes.append(box1)
|
211 |
+
else:
|
212 |
+
filtered_boxes.append(box1)
|
213 |
+
return torch.tensor(filtered_boxes)
|
214 |
+
|
215 |
+
def load_image(image_path: str) -> Tuple[np.array, torch.Tensor]:
|
216 |
+
transform = T.Compose(
|
217 |
+
[
|
218 |
+
T.RandomResize([800], max_size=1333),
|
219 |
+
T.ToTensor(),
|
220 |
+
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
221 |
+
]
|
222 |
+
)
|
223 |
+
image_source = Image.open(image_path).convert("RGB")
|
224 |
+
image = np.asarray(image_source)
|
225 |
+
image_transformed, _ = transform(image_source, None)
|
226 |
+
return image, image_transformed
|
227 |
+
|
228 |
+
|
229 |
+
def annotate(image_source: np.ndarray, boxes: torch.Tensor, logits: torch.Tensor, phrases: List[str], text_scale: float,
|
230 |
+
text_padding=5, text_thickness=2, thickness=3) -> np.ndarray:
|
231 |
+
"""
|
232 |
+
This function annotates an image with bounding boxes and labels.
|
233 |
+
|
234 |
+
Parameters:
|
235 |
+
image_source (np.ndarray): The source image to be annotated.
|
236 |
+
boxes (torch.Tensor): A tensor containing bounding box coordinates. in cxcywh format, pixel scale
|
237 |
+
logits (torch.Tensor): A tensor containing confidence scores for each bounding box.
|
238 |
+
phrases (List[str]): A list of labels for each bounding box.
|
239 |
+
text_scale (float): The scale of the text to be displayed. 0.8 for mobile/web, 0.3 for desktop # 0.4 for mind2web
|
240 |
+
|
241 |
+
Returns:
|
242 |
+
np.ndarray: The annotated image.
|
243 |
+
"""
|
244 |
+
h, w, _ = image_source.shape
|
245 |
+
boxes = boxes * torch.Tensor([w, h, w, h])
|
246 |
+
xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy()
|
247 |
+
xywh = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xywh").numpy()
|
248 |
+
detections = sv.Detections(xyxy=xyxy)
|
249 |
+
|
250 |
+
labels = [f"{phrase}" for phrase in range(boxes.shape[0])]
|
251 |
+
|
252 |
+
from util.box_annotator import BoxAnnotator
|
253 |
+
box_annotator = BoxAnnotator(text_scale=text_scale, text_padding=text_padding,text_thickness=text_thickness,thickness=thickness) # 0.8 for mobile/web, 0.3 for desktop # 0.4 for mind2web
|
254 |
+
annotated_frame = image_source.copy()
|
255 |
+
annotated_frame = box_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels, image_size=(w,h))
|
256 |
+
|
257 |
+
label_coordinates = {f"{phrase}": v for phrase, v in zip(phrases, xywh)}
|
258 |
+
return annotated_frame, label_coordinates
|
259 |
+
|
260 |
+
|
261 |
+
def predict(model, image, caption, box_threshold, text_threshold):
|
262 |
+
""" Use huggingface model to replace the original model
|
263 |
+
"""
|
264 |
+
model, processor = model['model'], model['processor']
|
265 |
+
device = model.device
|
266 |
+
|
267 |
+
inputs = processor(images=image, text=caption, return_tensors="pt").to(device)
|
268 |
+
with torch.no_grad():
|
269 |
+
outputs = model(**inputs)
|
270 |
+
|
271 |
+
results = processor.post_process_grounded_object_detection(
|
272 |
+
outputs,
|
273 |
+
inputs.input_ids,
|
274 |
+
box_threshold=box_threshold, # 0.4,
|
275 |
+
text_threshold=text_threshold, # 0.3,
|
276 |
+
target_sizes=[image.size[::-1]]
|
277 |
+
)[0]
|
278 |
+
boxes, logits, phrases = results["boxes"], results["scores"], results["labels"]
|
279 |
+
return boxes, logits, phrases
|
280 |
+
|
281 |
+
|
282 |
+
def predict_yolo(model, image_path, box_threshold):
|
283 |
+
""" Use huggingface model to replace the original model
|
284 |
+
"""
|
285 |
+
# model = model['model']
|
286 |
+
|
287 |
+
result = model.predict(
|
288 |
+
source=image_path,
|
289 |
+
conf=box_threshold,
|
290 |
+
# iou=0.5, # default 0.7
|
291 |
+
)
|
292 |
+
boxes = result[0].boxes.xyxy#.tolist() # in pixel space
|
293 |
+
conf = result[0].boxes.conf
|
294 |
+
phrases = [str(i) for i in range(len(boxes))]
|
295 |
+
|
296 |
+
return boxes, conf, phrases
|
297 |
+
|
298 |
+
|
299 |
+
def get_som_labeled_img(img_path, model=None, BOX_TRESHOLD = 0.01, output_coord_in_ratio=False, ocr_bbox=None, text_scale=0.4, text_padding=5, draw_bbox_config=None, caption_model_processor=None, ocr_text=[], use_local_semantics=True, iou_threshold=0.9,prompt=None):
|
300 |
+
""" ocr_bbox: list of xyxy format bbox
|
301 |
+
"""
|
302 |
+
TEXT_PROMPT = "clickable buttons on the screen"
|
303 |
+
# BOX_TRESHOLD = 0.02 # 0.05/0.02 for web and 0.1 for mobile
|
304 |
+
TEXT_TRESHOLD = 0.01 # 0.9 # 0.01
|
305 |
+
image_source = Image.open(img_path).convert("RGB")
|
306 |
+
w, h = image_source.size
|
307 |
+
# import pdb; pdb.set_trace()
|
308 |
+
if False: # TODO
|
309 |
+
xyxy, logits, phrases = predict(model=model, image=image_source, caption=TEXT_PROMPT, box_threshold=BOX_TRESHOLD, text_threshold=TEXT_TRESHOLD)
|
310 |
+
else:
|
311 |
+
xyxy, logits, phrases = predict_yolo(model=model, image_path=img_path, box_threshold=BOX_TRESHOLD)
|
312 |
+
xyxy = xyxy / torch.Tensor([w, h, w, h]).to(xyxy.device)
|
313 |
+
image_source = np.asarray(image_source)
|
314 |
+
phrases = [str(i) for i in range(len(phrases))]
|
315 |
+
|
316 |
+
# annotate the image with labels
|
317 |
+
h, w, _ = image_source.shape
|
318 |
+
if ocr_bbox:
|
319 |
+
ocr_bbox = torch.tensor(ocr_bbox) / torch.Tensor([w, h, w, h])
|
320 |
+
ocr_bbox=ocr_bbox.tolist()
|
321 |
+
else:
|
322 |
+
print('no ocr bbox!!!')
|
323 |
+
ocr_bbox = None
|
324 |
+
filtered_boxes = remove_overlap(boxes=xyxy, iou_threshold=iou_threshold, ocr_bbox=ocr_bbox)
|
325 |
+
|
326 |
+
# get parsed icon local semantics
|
327 |
+
if use_local_semantics:
|
328 |
+
caption_model = caption_model_processor['model']
|
329 |
+
if 'phi3_v' in caption_model.config.model_type:
|
330 |
+
parsed_content_icon = get_parsed_content_icon_phi3v(filtered_boxes, ocr_bbox, image_source, caption_model_processor)
|
331 |
+
else:
|
332 |
+
parsed_content_icon = get_parsed_content_icon(filtered_boxes, ocr_bbox, image_source, caption_model_processor, prompt=prompt)
|
333 |
+
ocr_text = [f"Text Box ID {i}: {txt}" for i, txt in enumerate(ocr_text)]
|
334 |
+
icon_start = len(ocr_text)
|
335 |
+
parsed_content_icon_ls = []
|
336 |
+
for i, txt in enumerate(parsed_content_icon):
|
337 |
+
parsed_content_icon_ls.append(f"Icon Box ID {str(i+icon_start)}: {txt}")
|
338 |
+
parsed_content_merged = ocr_text + parsed_content_icon_ls
|
339 |
+
else:
|
340 |
+
ocr_text = [f"Text Box ID {i}: {txt}" for i, txt in enumerate(ocr_text)]
|
341 |
+
parsed_content_merged = ocr_text
|
342 |
+
|
343 |
+
filtered_boxes = box_convert(boxes=filtered_boxes, in_fmt="xyxy", out_fmt="cxcywh")
|
344 |
+
|
345 |
+
phrases = [i for i in range(len(filtered_boxes))]
|
346 |
+
|
347 |
+
# draw boxes
|
348 |
+
if draw_bbox_config:
|
349 |
+
annotated_frame, label_coordinates = annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, **draw_bbox_config)
|
350 |
+
else:
|
351 |
+
annotated_frame, label_coordinates = annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, text_scale=text_scale, text_padding=text_padding)
|
352 |
+
|
353 |
+
pil_img = Image.fromarray(annotated_frame)
|
354 |
+
buffered = io.BytesIO()
|
355 |
+
pil_img.save(buffered, format="PNG")
|
356 |
+
encoded_image = base64.b64encode(buffered.getvalue()).decode('ascii')
|
357 |
+
if output_coord_in_ratio:
|
358 |
+
# h, w, _ = image_source.shape
|
359 |
+
label_coordinates = {k: [v[0]/w, v[1]/h, v[2]/w, v[3]/h] for k, v in label_coordinates.items()}
|
360 |
+
assert w == annotated_frame.shape[1] and h == annotated_frame.shape[0]
|
361 |
+
|
362 |
+
return encoded_image, label_coordinates, parsed_content_merged
|
363 |
+
|
364 |
+
|
365 |
+
def get_xywh(input):
|
366 |
+
x, y, w, h = input[0][0], input[0][1], input[2][0] - input[0][0], input[2][1] - input[0][1]
|
367 |
+
x, y, w, h = int(x), int(y), int(w), int(h)
|
368 |
+
return x, y, w, h
|
369 |
+
|
370 |
+
def get_xyxy(input):
|
371 |
+
x, y, xp, yp = input[0][0], input[0][1], input[2][0], input[2][1]
|
372 |
+
x, y, xp, yp = int(x), int(y), int(xp), int(yp)
|
373 |
+
return x, y, xp, yp
|
374 |
+
|
375 |
+
def get_xywh_yolo(input):
|
376 |
+
x, y, w, h = input[0], input[1], input[2] - input[0], input[3] - input[1]
|
377 |
+
x, y, w, h = int(x), int(y), int(w), int(h)
|
378 |
+
return x, y, w, h
|
379 |
+
|
380 |
+
|
381 |
+
|
382 |
+
def check_ocr_box(image_path, display_img = True, output_bb_format='xywh', goal_filtering=None, easyocr_args=None, use_paddleocr=False):
|
383 |
+
if use_paddleocr:
|
384 |
+
result = paddle_ocr.ocr(image_path, cls=False)[0]
|
385 |
+
coord = [item[0] for item in result]
|
386 |
+
text = [item[1][0] for item in result]
|
387 |
+
else: # EasyOCR
|
388 |
+
if easyocr_args is None:
|
389 |
+
easyocr_args = {}
|
390 |
+
result = reader.readtext(image_path, **easyocr_args)
|
391 |
+
# print('goal filtering pred:', result[-5:])
|
392 |
+
coord = [item[0] for item in result]
|
393 |
+
text = [item[1] for item in result]
|
394 |
+
# read the image using cv2
|
395 |
+
if display_img:
|
396 |
+
opencv_img = cv2.imread(image_path)
|
397 |
+
opencv_img = cv2.cvtColor(opencv_img, cv2.COLOR_RGB2BGR)
|
398 |
+
bb = []
|
399 |
+
for item in coord:
|
400 |
+
x, y, a, b = get_xywh(item)
|
401 |
+
# print(x, y, a, b)
|
402 |
+
bb.append((x, y, a, b))
|
403 |
+
cv2.rectangle(opencv_img, (x, y), (x+a, y+b), (0, 255, 0), 2)
|
404 |
+
|
405 |
+
# Display the image
|
406 |
+
plt.imshow(opencv_img)
|
407 |
+
else:
|
408 |
+
if output_bb_format == 'xywh':
|
409 |
+
bb = [get_xywh(item) for item in coord]
|
410 |
+
elif output_bb_format == 'xyxy':
|
411 |
+
bb = [get_xyxy(item) for item in coord]
|
412 |
+
# print('bounding box!!!', bb)
|
413 |
+
return (text, bb), goal_filtering
|
414 |
+
|
415 |
+
|
416 |
+
|