smthemex
@@ -0,0 +1,9 @@
|
||||
# Microsoft Open Source Code of Conduct
|
||||
|
||||
This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).
|
||||
|
||||
Resources:
|
||||
|
||||
- [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/)
|
||||
- [Microsoft Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/)
|
||||
- Contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with questions or concerns
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) Microsoft Corporation.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE
|
||||
@@ -0,0 +1,259 @@
|
||||
---
|
||||
license: mit
|
||||
license_link: https://huggingface.co/microsoft/Florence-2-base/resolve/main/LICENSE
|
||||
pipeline_tag: image-text-to-text
|
||||
tags:
|
||||
- vision
|
||||
---
|
||||
|
||||
# Florence-2: Advancing a Unified Representation for a Variety of Vision Tasks
|
||||
|
||||
## Model Summary
|
||||
|
||||
This Hub repository contains a HuggingFace's `transformers` implementation of Florence-2 model from Microsoft.
|
||||
|
||||
Florence-2 is an advanced vision foundation model that uses a prompt-based approach to handle a wide range of vision and vision-language tasks. Florence-2 can interpret simple text prompts to perform tasks like captioning, object detection, and segmentation. It leverages our FLD-5B dataset, containing 5.4 billion annotations across 126 million images, to master multi-task learning. The model's sequence-to-sequence architecture enables it to excel in both zero-shot and fine-tuned settings, proving to be a competitive vision foundation model.
|
||||
|
||||
Resources and Technical Documentation:
|
||||
+ [Florence-2 technical report](https://arxiv.org/abs/2311.06242).
|
||||
+ [Jupyter Notebook for inference and visualization of Florence-2-large model](https://huggingface.co/microsoft/Florence-2-large/blob/main/sample_inference.ipynb)
|
||||
|
||||
| Model | Model size | Model Description |
|
||||
| ------- | ------------- | ------------- |
|
||||
| Florence-2-base[[HF]](https://huggingface.co/microsoft/Florence-2-base) | 0.23B | Pretrained model with FLD-5B
|
||||
| Florence-2-large[[HF]](https://huggingface.co/microsoft/Florence-2-large) | 0.77B | Pretrained model with FLD-5B
|
||||
| Florence-2-base-ft[[HF]](https://huggingface.co/microsoft/Florence-2-base-ft) | 0.23B | Finetuned model on a colletion of downstream tasks
|
||||
| Florence-2-large-ft[[HF]](https://huggingface.co/microsoft/Florence-2-large-ft) | 0.77B | Finetuned model on a colletion of downstream tasks
|
||||
|
||||
## How to Get Started with the Model
|
||||
|
||||
Use the code below to get started with the model. All models are trained with float16.
|
||||
|
||||
```python
|
||||
import requests
|
||||
|
||||
from PIL import Image
|
||||
from transformers import AutoProcessor, AutoModelForCausalLM
|
||||
|
||||
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||||
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-base", torch_dtype=torch_dtype, trust_remote_code=True).to(device)
|
||||
processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
|
||||
|
||||
prompt = "<OD>"
|
||||
|
||||
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true"
|
||||
image = Image.open(requests.get(url, stream=True).raw)
|
||||
|
||||
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=1024,
|
||||
do_sample=False,
|
||||
num_beams=3,
|
||||
)
|
||||
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
|
||||
|
||||
parsed_answer = processor.post_process_generation(generated_text, task="<OD>", image_size=(image.width, image.height))
|
||||
|
||||
print(parsed_answer)
|
||||
|
||||
```
|
||||
|
||||
|
||||
## Tasks
|
||||
|
||||
This model is capable of performing different tasks through changing the prompts.
|
||||
|
||||
First, let's define a function to run a prompt.
|
||||
|
||||
<details>
|
||||
<summary> Click to expand </summary>
|
||||
|
||||
```python
|
||||
import requests
|
||||
|
||||
from PIL import Image
|
||||
from transformers import AutoProcessor, AutoModelForCausalLM
|
||||
|
||||
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||||
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-base", torch_dtype=torch_dtype, trust_remote_code=True).to(device)
|
||||
processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
|
||||
|
||||
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true"
|
||||
image = Image.open(requests.get(url, stream=True).raw)
|
||||
|
||||
def run_example(task_prompt, text_input=None):
|
||||
if text_input is None:
|
||||
prompt = task_prompt
|
||||
else:
|
||||
prompt = task_prompt + text_input
|
||||
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=1024,
|
||||
num_beams=3
|
||||
)
|
||||
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
|
||||
|
||||
parsed_answer = processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
|
||||
|
||||
print(parsed_answer)
|
||||
```
|
||||
</details>
|
||||
|
||||
Here are the tasks `Florence-2` could perform:
|
||||
|
||||
<details>
|
||||
<summary> Click to expand </summary>
|
||||
|
||||
|
||||
|
||||
### Caption
|
||||
```python
|
||||
prompt = "<CAPTION>"
|
||||
run_example(prompt)
|
||||
```
|
||||
|
||||
### Detailed Caption
|
||||
```python
|
||||
prompt = "<DETAILED_CAPTION>"
|
||||
run_example(prompt)
|
||||
```
|
||||
|
||||
### More Detailed Caption
|
||||
```python
|
||||
prompt = "<MORE_DETAILED_CAPTION>"
|
||||
run_example(prompt)
|
||||
```
|
||||
|
||||
### Caption to Phrase Grounding
|
||||
caption to phrase grounding task requires additional text input, i.e. caption.
|
||||
|
||||
Caption to phrase grounding results format:
|
||||
{'\<CAPTION_TO_PHRASE_GROUNDING>': {'bboxes': [[x1, y1, x2, y2], ...], 'labels': ['', '', ...]}}
|
||||
```python
|
||||
task_prompt = "<CAPTION_TO_PHRASE_GROUNDING>"
|
||||
results = run_example(task_prompt, text_input="A green car parked in front of a yellow building.")
|
||||
```
|
||||
|
||||
### Object Detection
|
||||
|
||||
OD results format:
|
||||
{'\<OD>': {'bboxes': [[x1, y1, x2, y2], ...],
|
||||
'labels': ['label1', 'label2', ...]} }
|
||||
|
||||
```python
|
||||
prompt = "<OD>"
|
||||
run_example(prompt)
|
||||
```
|
||||
|
||||
### Dense Region Caption
|
||||
Dense region caption results format:
|
||||
{'\<DENSE_REGION_CAPTION>' : {'bboxes': [[x1, y1, x2, y2], ...],
|
||||
'labels': ['label1', 'label2', ...]} }
|
||||
```python
|
||||
prompt = "<DENSE_REGION_CAPTION>"
|
||||
run_example(prompt)
|
||||
```
|
||||
|
||||
### Region proposal
|
||||
Dense region caption results format:
|
||||
{'\<REGION_PROPOSAL>': {'bboxes': [[x1, y1, x2, y2], ...],
|
||||
'labels': ['', '', ...]}}
|
||||
```python
|
||||
prompt = "<REGION_PROPOSAL>"
|
||||
run_example(prompt)
|
||||
```
|
||||
|
||||
### OCR
|
||||
|
||||
```python
|
||||
prompt = "<OCR>"
|
||||
run_example(prompt)
|
||||
```
|
||||
|
||||
### OCR with Region
|
||||
OCR with region output format:
|
||||
{'\<OCR_WITH_REGION>': {'quad_boxes': [[x1, y1, x2, y2, x3, y3, x4, y4], ...], 'labels': ['text1', ...]}}
|
||||
```python
|
||||
prompt = "<OCR_WITH_REGION>"
|
||||
run_example(prompt)
|
||||
```
|
||||
|
||||
for More detailed examples, please refer to [notebook](https://huggingface.co/microsoft/Florence-2-large/blob/main/sample_inference.ipynb)
|
||||
</details>
|
||||
|
||||
# Benchmarks
|
||||
|
||||
## Florence-2 Zero-shot performance
|
||||
|
||||
The following table presents the zero-shot performance of generalist vision foundation models on image captioning and object detection evaluation tasks. These models have not been exposed to the training data of the evaluation tasks during their training phase.
|
||||
|
||||
| Method | #params | COCO Cap. test CIDEr | NoCaps val CIDEr | TextCaps val CIDEr | COCO Det. val2017 mAP |
|
||||
|--------|---------|----------------------|------------------|--------------------|-----------------------|
|
||||
| Flamingo | 80B | 84.3 | - | - | - |
|
||||
| Florence-2-base| 0.23B | 133.0 | 118.7 | 70.1 | 34.7 |
|
||||
| Florence-2-large| 0.77B | 135.6 | 120.8 | 72.8 | 37.5 |
|
||||
|
||||
|
||||
The following table continues the comparison with performance on other vision-language evaluation tasks.
|
||||
|
||||
| Method | Flickr30k test R@1 | Refcoco val Accuracy | Refcoco test-A Accuracy | Refcoco test-B Accuracy | Refcoco+ val Accuracy | Refcoco+ test-A Accuracy | Refcoco+ test-B Accuracy | Refcocog val Accuracy | Refcocog test Accuracy | Refcoco RES val mIoU |
|
||||
|--------|----------------------|----------------------|-------------------------|-------------------------|-----------------------|--------------------------|--------------------------|-----------------------|------------------------|----------------------|
|
||||
| Kosmos-2 | 78.7 | 52.3 | 57.4 | 47.3 | 45.5 | 50.7 | 42.2 | 60.6 | 61.7 | - |
|
||||
| Florence-2-base | 83.6 | 53.9 | 58.4 | 49.7 | 51.5 | 56.4 | 47.9 | 66.3 | 65.1 | 34.6 |
|
||||
| Florence-2-large | 84.4 | 56.3 | 61.6 | 51.4 | 53.6 | 57.9 | 49.9 | 68.0 | 67.0 | 35.8 |
|
||||
|
||||
|
||||
|
||||
## Florence-2 finetuned performance
|
||||
|
||||
We finetune Florence-2 models with a collection of downstream tasks, resulting two generalist models *Florence-2-base-ft* and *Florence-2-large-ft* that can conduct a wide range of downstream tasks.
|
||||
|
||||
The table below compares the performance of specialist and generalist models on various captioning and Visual Question Answering (VQA) tasks. Specialist models are fine-tuned specifically for each task, whereas generalist models are fine-tuned in a task-agnostic manner across all tasks. The symbol "▲" indicates the usage of external OCR as input.
|
||||
|
||||
| Method | # Params | COCO Caption Karpathy test CIDEr | NoCaps val CIDEr | TextCaps val CIDEr | VQAv2 test-dev Acc | TextVQA test-dev Acc | VizWiz VQA test-dev Acc |
|
||||
|----------------|----------|-----------------------------------|------------------|--------------------|--------------------|----------------------|-------------------------|
|
||||
| **Specialist Models** | | | | | | | |
|
||||
| CoCa | 2.1B | 143.6 | 122.4 | - | 82.3 | - | - |
|
||||
| BLIP-2 | 7.8B | 144.5 | 121.6 | - | 82.2 | - | - |
|
||||
| GIT2 | 5.1B | 145.0 | 126.9 | 148.6 | 81.7 | 67.3 | 71.0 |
|
||||
| Flamingo | 80B | 138.1 | - | - | 82.0 | 54.1 | 65.7 |
|
||||
| PaLI | 17B | 149.1 | 127.0 | 160.0▲ | 84.3 | 58.8 / 73.1▲ | 71.6 / 74.4▲ |
|
||||
| PaLI-X | 55B | 149.2 | 126.3 | 147.0 / 163.7▲ | 86.0 | 71.4 / 80.8▲ | 70.9 / 74.6▲ |
|
||||
| **Generalist Models** | | | | | | | |
|
||||
| Unified-IO | 2.9B | - | 100.0 | - | 77.9 | - | 57.4 |
|
||||
| Florence-2-base-ft | 0.23B | 140.0 | 116.7 | 143.9 | 79.7 | 63.6 | 63.6 |
|
||||
| Florence-2-large-ft | 0.77B | 143.3 | 124.9 | 151.1 | 81.7 | 73.5 | 72.6 |
|
||||
|
||||
|
||||
| Method | # Params | COCO Det. val2017 mAP | Flickr30k test R@1 | RefCOCO val Accuracy | RefCOCO test-A Accuracy | RefCOCO test-B Accuracy | RefCOCO+ val Accuracy | RefCOCO+ test-A Accuracy | RefCOCO+ test-B Accuracy | RefCOCOg val Accuracy | RefCOCOg test Accuracy | RefCOCO RES val mIoU |
|
||||
|----------------------|----------|-----------------------|--------------------|----------------------|-------------------------|-------------------------|------------------------|---------------------------|---------------------------|------------------------|-----------------------|------------------------|
|
||||
| **Specialist Models** | | | | | | | | | | | | |
|
||||
| SeqTR | - | - | - | 83.7 | 86.5 | 81.2 | 71.5 | 76.3 | 64.9 | 74.9 | 74.2 | - |
|
||||
| PolyFormer | - | - | - | 90.4 | 92.9 | 87.2 | 85.0 | 89.8 | 78.0 | 85.8 | 85.9 | 76.9 |
|
||||
| UNINEXT | 0.74B | 60.6 | - | 92.6 | 94.3 | 91.5 | 85.2 | 89.6 | 79.8 | 88.7 | 89.4 | - |
|
||||
| Ferret | 13B | - | - | 89.5 | 92.4 | 84.4 | 82.8 | 88.1 | 75.2 | 85.8 | 86.3 | - |
|
||||
| **Generalist Models** | | | | | | | | | | | | |
|
||||
| UniTAB | - | - | - | 88.6 | 91.1 | 83.8 | 81.0 | 85.4 | 71.6 | 84.6 | 84.7 | - |
|
||||
| Florence-2-base-ft | 0.23B | 41.4 | 84.0 | 92.6 | 94.8 | 91.5 | 86.8 | 91.7 | 82.2 | 89.8 | 82.2 | 78.0 |
|
||||
| Florence-2-large-ft| 0.77B | 43.4 | 85.2 | 93.4 | 95.3 | 92.0 | 88.3 | 92.9 | 83.6 | 91.2 | 91.7 | 80.5 |
|
||||
|
||||
|
||||
## BibTex and citation info
|
||||
|
||||
```
|
||||
@article{xiao2023florence,
|
||||
title={Florence-2: Advancing a unified representation for a variety of vision tasks},
|
||||
author={Xiao, Bin and Wu, Haiping and Xu, Weijian and Dai, Xiyang and Hu, Houdong and Lu, Yumao and Zeng, Michael and Liu, Ce and Yuan, Lu},
|
||||
journal={arXiv preprint arXiv:2311.06242},
|
||||
year={2023}
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,41 @@
|
||||
<!-- BEGIN MICROSOFT SECURITY.MD V0.0.9 BLOCK -->
|
||||
|
||||
## Security
|
||||
|
||||
Microsoft takes the security of our software products and services seriously, which includes all source code repositories managed through our GitHub organizations, which include [Microsoft](https://github.com/Microsoft), [Azure](https://github.com/Azure), [DotNet](https://github.com/dotnet), [AspNet](https://github.com/aspnet) and [Xamarin](https://github.com/xamarin).
|
||||
|
||||
If you believe you have found a security vulnerability in any Microsoft-owned repository that meets [Microsoft's definition of a security vulnerability](https://aka.ms/security.md/definition), please report it to us as described below.
|
||||
|
||||
## Reporting Security Issues
|
||||
|
||||
**Please do not report security vulnerabilities through public GitHub issues.**
|
||||
|
||||
Instead, please report them to the Microsoft Security Response Center (MSRC) at [https://msrc.microsoft.com/create-report](https://aka.ms/security.md/msrc/create-report).
|
||||
|
||||
If you prefer to submit without logging in, send email to [secure@microsoft.com](mailto:secure@microsoft.com). If possible, encrypt your message with our PGP key; please download it from the [Microsoft Security Response Center PGP Key page](https://aka.ms/security.md/msrc/pgp).
|
||||
|
||||
You should receive a response within 24 hours. If for some reason you do not, please follow up via email to ensure we received your original message. Additional information can be found at [microsoft.com/msrc](https://www.microsoft.com/msrc).
|
||||
|
||||
Please include the requested information listed below (as much as you can provide) to help us better understand the nature and scope of the possible issue:
|
||||
|
||||
* Type of issue (e.g. buffer overflow, SQL injection, cross-site scripting, etc.)
|
||||
* Full paths of source file(s) related to the manifestation of the issue
|
||||
* The location of the affected source code (tag/branch/commit or direct URL)
|
||||
* Any special configuration required to reproduce the issue
|
||||
* Step-by-step instructions to reproduce the issue
|
||||
* Proof-of-concept or exploit code (if possible)
|
||||
* Impact of the issue, including how an attacker might exploit the issue
|
||||
|
||||
This information will help us triage your report more quickly.
|
||||
|
||||
If you are reporting for a bug bounty, more complete reports can contribute to a higher bounty award. Please visit our [Microsoft Bug Bounty Program](https://aka.ms/security.md/msrc/bounty) page for more details about our active programs.
|
||||
|
||||
## Preferred Languages
|
||||
|
||||
We prefer all communications to be in English.
|
||||
|
||||
## Policy
|
||||
|
||||
Microsoft follows the principle of [Coordinated Vulnerability Disclosure](https://aka.ms/security.md/cvd).
|
||||
|
||||
<!-- END MICROSOFT SECURITY.MD BLOCK -->
|
||||
@@ -0,0 +1,25 @@
|
||||
# TODO: The maintainer of this repo has not yet edited this file
|
||||
|
||||
**REPO OWNER**: Do you want Customer Service & Support (CSS) support for this product/project?
|
||||
|
||||
- **No CSS support:** Fill out this template with information about how to file issues and get help.
|
||||
- **Yes CSS support:** Fill out an intake form at [aka.ms/onboardsupport](https://aka.ms/onboardsupport). CSS will work with/help you to determine next steps.
|
||||
- **Not sure?** Fill out an intake as though the answer were "Yes". CSS will help you decide.
|
||||
|
||||
*Then remove this first heading from this SUPPORT.MD file before publishing your repo.*
|
||||
|
||||
# Support
|
||||
|
||||
## How to file issues and get help
|
||||
|
||||
This project uses GitHub Issues to track bugs and feature requests. Please search the existing
|
||||
issues before filing new issues to avoid duplicates. For new issues, file your bug or
|
||||
feature request as a new Issue.
|
||||
|
||||
For help and questions about using this project, please **REPO MAINTAINER: INSERT INSTRUCTIONS HERE
|
||||
FOR HOW TO ENGAGE REPO OWNERS OR COMMUNITY FOR HELP. COULD BE A STACK OVERFLOW TAG OR OTHER
|
||||
CHANNEL. WHERE WILL YOU HELP PEOPLE?**.
|
||||
|
||||
## Microsoft Support Policy
|
||||
|
||||
Support for this **PROJECT or PRODUCT** is limited to the resources listed above.
|
||||
@@ -0,0 +1,85 @@
|
||||
{
|
||||
"_name_or_path": "florence2",
|
||||
"architectures": [
|
||||
"Florence2ForConditionalGeneration"
|
||||
],
|
||||
"auto_map": {
|
||||
"AutoConfig": "configuration_florence2.Florence2Config",
|
||||
"AutoModelForCausalLM": "modeling_florence2.Florence2ForConditionalGeneration"
|
||||
},
|
||||
"bos_token_id": 0,
|
||||
"eos_token_id": 2,
|
||||
"ignore_index": -100,
|
||||
"model_type": "florence2",
|
||||
"pad_token_id": 1,
|
||||
"projection_dim": 768,
|
||||
"text_config": {
|
||||
"vocab_size": 51289,
|
||||
"activation_dropout": 0.1,
|
||||
"activation_function": "gelu",
|
||||
"add_bias_logits": false,
|
||||
"add_final_layer_norm": false,
|
||||
"attention_dropout": 0.1,
|
||||
"bos_token_id": 0,
|
||||
"classif_dropout": 0.1,
|
||||
"classifier_dropout": 0.0,
|
||||
"d_model": 768,
|
||||
"decoder_attention_heads": 12,
|
||||
"decoder_ffn_dim": 3072,
|
||||
"decoder_layerdrop": 0.0,
|
||||
"decoder_layers": 6,
|
||||
"decoder_start_token_id": 2,
|
||||
"dropout": 0.1,
|
||||
"early_stopping": true,
|
||||
"encoder_attention_heads": 12,
|
||||
"encoder_ffn_dim": 3072,
|
||||
"encoder_layerdrop": 0.0,
|
||||
"encoder_layers": 6,
|
||||
"eos_token_id": 2,
|
||||
"forced_eos_token_id": 2,
|
||||
"forced_bos_token_id": 0,
|
||||
"gradient_checkpointing": false,
|
||||
"init_std": 0.02,
|
||||
"is_encoder_decoder": true,
|
||||
"label2id": {
|
||||
"LABEL_0": 0,
|
||||
"LABEL_1": 1,
|
||||
"LABEL_2": 2
|
||||
},
|
||||
"max_position_embeddings": 1024,
|
||||
"no_repeat_ngram_size": 3,
|
||||
"normalize_before": false,
|
||||
"num_hidden_layers": 6,
|
||||
"pad_token_id": 1,
|
||||
"scale_embedding": false,
|
||||
"num_beams": 3
|
||||
},
|
||||
"vision_config": {
|
||||
"model_type": "davit",
|
||||
"drop_path_rate": 0.1,
|
||||
"patch_size": [7, 3, 3, 3],
|
||||
"patch_stride": [4, 2, 2, 2],
|
||||
"patch_padding": [3, 1, 1, 1],
|
||||
"patch_prenorm": [false, true, true, true],
|
||||
"enable_checkpoint": false,
|
||||
"dim_embed": [128, 256, 512, 1024],
|
||||
"num_heads": [4, 8, 16, 32],
|
||||
"num_groups": [4, 8, 16, 32],
|
||||
"depths": [1, 1, 9, 1],
|
||||
"window_size": 12,
|
||||
"projection_dim": 768,
|
||||
"visual_temporal_embedding": {
|
||||
"type": "COSINE",
|
||||
"max_temporal_embeddings": 100
|
||||
},
|
||||
"image_pos_embed": {
|
||||
"type": "learned_abs_2d",
|
||||
"max_pos_embeddings": 50
|
||||
},
|
||||
"image_feature_source": ["spatial_avg_pool", "temporal_avg_pool"]
|
||||
},
|
||||
"vocab_size": 51289,
|
||||
"torch_dtype": "float16",
|
||||
"transformers_version": "4.41.0.dev0",
|
||||
"is_encoder_decoder": true
|
||||
}
|
||||
@@ -0,0 +1,340 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2024 Microsoft and the HuggingFace Inc. team. All rights reserved.
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
import warnings
|
||||
""" Florence-2 configuration"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from transformers import AutoConfig
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.utils import logging
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
class Florence2VisionConfig(PretrainedConfig):
|
||||
r"""
|
||||
This is the configuration class to store the configuration of a [`Florence2VisionModel`]. It is used to instantiate a Florence2VisionModel
|
||||
according to the specified arguments, defining the model architecture. Instantiating a configuration with the
|
||||
defaults will yield a similar configuration to that of the Florence2VisionModel architecture.
|
||||
|
||||
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||
documentation from [`PretrainedConfig`] for more information.
|
||||
|
||||
Args:
|
||||
drop_path_rate (`float`, *optional*, defaults to 0.1):
|
||||
The dropout rate of the drop path layer.
|
||||
patch_size (`List[int]`, *optional*, defaults to [7, 3, 3, 3]):
|
||||
The patch size of the image.
|
||||
patch_stride (`List[int]`, *optional*, defaults to [4, 2, 2, 2]):
|
||||
The patch stride of the image.
|
||||
patch_padding (`List[int]`, *optional*, defaults to [3, 1, 1, 1]):
|
||||
The patch padding of the image.
|
||||
patch_prenorm (`List[bool]`, *optional*, defaults to [false, true, true, true]):
|
||||
Whether to apply layer normalization before the patch embedding layer.
|
||||
enable_checkpoint (`bool`, *optional*, defaults to False):
|
||||
Whether to enable checkpointing.
|
||||
dim_embed (`List[int]`, *optional*, defaults to [256, 512, 1024, 2048]):
|
||||
The dimension of the embedding layer.
|
||||
num_heads (`List[int]`, *optional*, defaults to [8, 16, 32, 64]):
|
||||
The number of attention heads.
|
||||
num_groups (`List[int]`, *optional*, defaults to [8, 16, 32, 64]):
|
||||
The number of groups.
|
||||
depths (`List[int]`, *optional*, defaults to [1, 1, 9, 1]):
|
||||
The depth of the model.
|
||||
window_size (`int`, *optional*, defaults to 12):
|
||||
The window size of the model.
|
||||
projection_dim (`int`, *optional*, defaults to 1024):
|
||||
The dimension of the projection layer.
|
||||
visual_temporal_embedding (`dict`, *optional*):
|
||||
The configuration of the visual temporal embedding.
|
||||
image_pos_embed (`dict`, *optional*):
|
||||
The configuration of the image position embedding.
|
||||
image_feature_source (`List[str]`, *optional*, defaults to ["spatial_avg_pool", "temporal_avg_pool"]):
|
||||
The source of the image feature.
|
||||
Example:
|
||||
|
||||
```python
|
||||
>>> from transformers import Florence2VisionConfig, Florence2VisionModel
|
||||
|
||||
>>> # Initializing a Florence2 Vision style configuration
|
||||
>>> configuration = Florence2VisionConfig()
|
||||
|
||||
>>> # Initializing a model (with random weights)
|
||||
>>> model = Florence2VisionModel(configuration)
|
||||
|
||||
>>> # Accessing the model configuration
|
||||
>>> configuration = model.config
|
||||
```"""
|
||||
|
||||
model_type = "florence2_vision"
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
drop_path_rate=0.1,
|
||||
patch_size=[7, 3, 3, 3],
|
||||
patch_stride=[4, 2, 2, 2],
|
||||
patch_padding=[3, 1, 1, 1],
|
||||
patch_prenorm=[False, True, True, True],
|
||||
enable_checkpoint=False,
|
||||
dim_embed=[256, 512, 1024, 2048],
|
||||
num_heads=[8, 16, 32, 64],
|
||||
num_groups=[8, 16, 32, 64],
|
||||
depths=[1, 1, 9, 1],
|
||||
window_size=12,
|
||||
projection_dim=1024,
|
||||
visual_temporal_embedding=None,
|
||||
image_pos_embed=None,
|
||||
image_feature_source=["spatial_avg_pool", "temporal_avg_pool"],
|
||||
**kwargs,
|
||||
):
|
||||
self.drop_path_rate = drop_path_rate
|
||||
self.patch_size = patch_size
|
||||
self.patch_stride = patch_stride
|
||||
self.patch_padding = patch_padding
|
||||
self.patch_prenorm = patch_prenorm
|
||||
self.enable_checkpoint = enable_checkpoint
|
||||
self.dim_embed = dim_embed
|
||||
self.num_heads = num_heads
|
||||
self.num_groups = num_groups
|
||||
self.depths = depths
|
||||
self.window_size = window_size
|
||||
self.projection_dim = projection_dim
|
||||
self.visual_temporal_embedding = visual_temporal_embedding
|
||||
self.image_pos_embed = image_pos_embed
|
||||
self.image_feature_source = image_feature_source
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
|
||||
|
||||
class Florence2LanguageConfig(PretrainedConfig):
|
||||
r"""
|
||||
This is the configuration class to store the configuration of a [`Florence2LanguagePreTrainedModel`]. It is used to instantiate a BART
|
||||
model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
|
||||
defaults will yield a similar configuration to that of the BART
|
||||
[facebook/bart-large](https://huggingface.co/facebook/bart-large) architecture.
|
||||
|
||||
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||
documentation from [`PretrainedConfig`] for more information.
|
||||
|
||||
|
||||
Args:
|
||||
vocab_size (`int`, *optional*, defaults to 51289):
|
||||
Vocabulary size of the Florence2Language model. Defines the number of different tokens that can be represented by the
|
||||
`inputs_ids` passed when calling [`Florence2LanguageModel`].
|
||||
d_model (`int`, *optional*, defaults to 1024):
|
||||
Dimensionality of the layers and the pooler layer.
|
||||
encoder_layers (`int`, *optional*, defaults to 12):
|
||||
Number of encoder layers.
|
||||
decoder_layers (`int`, *optional*, defaults to 12):
|
||||
Number of decoder layers.
|
||||
encoder_attention_heads (`int`, *optional*, defaults to 16):
|
||||
Number of attention heads for each attention layer in the Transformer encoder.
|
||||
decoder_attention_heads (`int`, *optional*, defaults to 16):
|
||||
Number of attention heads for each attention layer in the Transformer decoder.
|
||||
decoder_ffn_dim (`int`, *optional*, defaults to 4096):
|
||||
Dimensionality of the "intermediate" (often named feed-forward) layer in decoder.
|
||||
encoder_ffn_dim (`int`, *optional*, defaults to 4096):
|
||||
Dimensionality of the "intermediate" (often named feed-forward) layer in decoder.
|
||||
activation_function (`str` or `function`, *optional*, defaults to `"gelu"`):
|
||||
The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
|
||||
`"relu"`, `"silu"` and `"gelu_new"` are supported.
|
||||
dropout (`float`, *optional*, defaults to 0.1):
|
||||
The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
|
||||
attention_dropout (`float`, *optional*, defaults to 0.0):
|
||||
The dropout ratio for the attention probabilities.
|
||||
activation_dropout (`float`, *optional*, defaults to 0.0):
|
||||
The dropout ratio for activations inside the fully connected layer.
|
||||
classifier_dropout (`float`, *optional*, defaults to 0.0):
|
||||
The dropout ratio for classifier.
|
||||
max_position_embeddings (`int`, *optional*, defaults to 1024):
|
||||
The maximum sequence length that this model might ever be used with. Typically set this to something large
|
||||
just in case (e.g., 512 or 1024 or 2048).
|
||||
init_std (`float`, *optional*, defaults to 0.02):
|
||||
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
|
||||
encoder_layerdrop (`float`, *optional*, defaults to 0.0):
|
||||
The LayerDrop probability for the encoder. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556)
|
||||
for more details.
|
||||
decoder_layerdrop (`float`, *optional*, defaults to 0.0):
|
||||
The LayerDrop probability for the decoder. See the [LayerDrop paper](see https://arxiv.org/abs/1909.11556)
|
||||
for more details.
|
||||
scale_embedding (`bool`, *optional*, defaults to `False`):
|
||||
Scale embeddings by diving by sqrt(d_model).
|
||||
use_cache (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not the model should return the last key/values attentions (not used by all models).
|
||||
num_labels (`int`, *optional*, defaults to 3):
|
||||
The number of labels to use in [`Florence2LanguageForSequenceClassification`].
|
||||
forced_eos_token_id (`int`, *optional*, defaults to 2):
|
||||
The id of the token to force as the last generated token when `max_length` is reached. Usually set to
|
||||
`eos_token_id`.
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
>>> from transformers import Florence2LanguageConfig, Florence2LanguageModel
|
||||
|
||||
>>> # Initializing a Florence2 Language style configuration
|
||||
>>> configuration = Florence2LanguageConfig()
|
||||
|
||||
>>> # Initializing a model (with random weights)
|
||||
>>> model = Florence2LangaugeModel(configuration)
|
||||
|
||||
>>> # Accessing the model configuration
|
||||
>>> configuration = model.config
|
||||
```"""
|
||||
|
||||
model_type = "florence2_language"
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
attribute_map = {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size=51289,
|
||||
max_position_embeddings=1024,
|
||||
encoder_layers=12,
|
||||
encoder_ffn_dim=4096,
|
||||
encoder_attention_heads=16,
|
||||
decoder_layers=12,
|
||||
decoder_ffn_dim=4096,
|
||||
decoder_attention_heads=16,
|
||||
encoder_layerdrop=0.0,
|
||||
decoder_layerdrop=0.0,
|
||||
activation_function="gelu",
|
||||
d_model=1024,
|
||||
dropout=0.1,
|
||||
attention_dropout=0.0,
|
||||
activation_dropout=0.0,
|
||||
init_std=0.02,
|
||||
classifier_dropout=0.0,
|
||||
scale_embedding=False,
|
||||
use_cache=True,
|
||||
num_labels=3,
|
||||
pad_token_id=1,
|
||||
bos_token_id=0,
|
||||
eos_token_id=2,
|
||||
is_encoder_decoder=True,
|
||||
decoder_start_token_id=2,
|
||||
forced_eos_token_id=2,
|
||||
**kwargs,
|
||||
):
|
||||
self.vocab_size = vocab_size
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.d_model = d_model
|
||||
self.encoder_ffn_dim = encoder_ffn_dim
|
||||
self.encoder_layers = encoder_layers
|
||||
self.encoder_attention_heads = encoder_attention_heads
|
||||
self.decoder_ffn_dim = decoder_ffn_dim
|
||||
self.decoder_layers = decoder_layers
|
||||
self.decoder_attention_heads = decoder_attention_heads
|
||||
self.dropout = dropout
|
||||
self.attention_dropout = attention_dropout
|
||||
self.activation_dropout = activation_dropout
|
||||
self.activation_function = activation_function
|
||||
self.init_std = init_std
|
||||
self.encoder_layerdrop = encoder_layerdrop
|
||||
self.decoder_layerdrop = decoder_layerdrop
|
||||
self.classifier_dropout = classifier_dropout
|
||||
self.use_cache = use_cache
|
||||
self.num_hidden_layers = encoder_layers
|
||||
self.scale_embedding = scale_embedding # scale factor will be sqrt(d_model) if True
|
||||
|
||||
super().__init__(
|
||||
num_labels=num_labels,
|
||||
pad_token_id=pad_token_id,
|
||||
bos_token_id=bos_token_id,
|
||||
eos_token_id=eos_token_id,
|
||||
is_encoder_decoder=is_encoder_decoder,
|
||||
decoder_start_token_id=decoder_start_token_id,
|
||||
forced_eos_token_id=forced_eos_token_id,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# ensure backward compatibility for BART CNN models
|
||||
if self.forced_bos_token_id is None and kwargs.get("force_bos_token_to_be_generated", False):
|
||||
self.forced_bos_token_id = self.bos_token_id
|
||||
warnings.warn(
|
||||
f"Please make sure the config includes `forced_bos_token_id={self.bos_token_id}` in future versions. "
|
||||
"The config can simply be saved and uploaded again to be fixed."
|
||||
)
|
||||
|
||||
class Florence2Config(PretrainedConfig):
|
||||
r"""
|
||||
This is the configuration class to store the configuration of a [`Florence2ForConditionalGeneration`]. It is used to instantiate an
|
||||
Florence-2 model according to the specified arguments, defining the model architecture.
|
||||
|
||||
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||
documentation from [`PretrainedConfig`] for more information.
|
||||
|
||||
Args:
|
||||
vision_config (`Florence2VisionConfig`, *optional*):
|
||||
Custom vision config or dict
|
||||
text_config (`Union[AutoConfig, dict]`, *optional*):
|
||||
The config object of the text backbone.
|
||||
ignore_index (`int`, *optional*, defaults to -100):
|
||||
The ignore index for the loss function.
|
||||
vocab_size (`int`, *optional*, defaults to 51289):
|
||||
Vocabulary size of the Florence2model. Defines the number of different tokens that can be represented by the
|
||||
`inputs_ids` passed when calling [`~Florence2ForConditionalGeneration`]
|
||||
projection_dim (`int`, *optional*, defaults to 1024):
|
||||
Dimension of the multimodal projection space.
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
>>> from transformers import Florence2ForConditionalGeneration, Florence2Config, CLIPVisionConfig, BartConfig
|
||||
|
||||
>>> # Initializing a clip-like vision config
|
||||
>>> vision_config = CLIPVisionConfig()
|
||||
|
||||
>>> # Initializing a Bart config
|
||||
>>> text_config = BartConfig()
|
||||
|
||||
>>> # Initializing a Florence-2 configuration
|
||||
>>> configuration = Florence2Config(vision_config, text_config)
|
||||
|
||||
>>> # Initializing a model from the florence-2 configuration
|
||||
>>> model = Florence2ForConditionalGeneration(configuration)
|
||||
|
||||
>>> # Accessing the model configuration
|
||||
>>> configuration = model.config
|
||||
```"""
|
||||
|
||||
model_type = "florence2"
|
||||
is_composition = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vision_config=None,
|
||||
text_config=None,
|
||||
ignore_index=-100,
|
||||
vocab_size=51289,
|
||||
projection_dim=1024,
|
||||
**kwargs,
|
||||
):
|
||||
self.ignore_index = ignore_index
|
||||
self.vocab_size = vocab_size
|
||||
self.projection_dim = projection_dim
|
||||
if vision_config is not None:
|
||||
vision_config = PretrainedConfig(**vision_config)
|
||||
self.vision_config = vision_config
|
||||
self.vocab_size = self.vocab_size
|
||||
|
||||
self.text_config = text_config
|
||||
if text_config is not None:
|
||||
self.text_config = Florence2LanguageConfig(**text_config)
|
||||
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,39 @@
|
||||
{
|
||||
"auto_map": {
|
||||
"AutoProcessor": "processing_florence2.Florence2Processor"
|
||||
},
|
||||
"_valid_processor_keys": [
|
||||
"images",
|
||||
"do_resize",
|
||||
"size",
|
||||
"resample",
|
||||
"do_rescale",
|
||||
"rescale_factor",
|
||||
"do_normalize",
|
||||
"image_mean",
|
||||
"image_std",
|
||||
"return_tensors",
|
||||
"data_format",
|
||||
"input_data_format",
|
||||
"do_convert_rgb"
|
||||
],
|
||||
"do_convert_rgb": null,
|
||||
"do_normalize": true,
|
||||
"do_rescale": true,
|
||||
"do_resize": true,
|
||||
"do_center_crop": false,
|
||||
"image_processor_type": "CLIPImageProcessor",
|
||||
"image_seq_length": 577,
|
||||
"image_mean": [0.485, 0.456, 0.406],
|
||||
"image_std": [0.229, 0.224, 0.225],
|
||||
"processor_class": "Florence2Processor",
|
||||
"resample": 3,
|
||||
"size": {
|
||||
"height": 768,
|
||||
"width":768
|
||||
},
|
||||
"crop_size": {
|
||||
"height": 768,
|
||||
"width": 768
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because one or more lines are too long
@@ -0,0 +1,4 @@
|
||||
{
|
||||
"model_max_length": 1024
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,56 @@
|
||||
# OmniParser: Screen Parsing tool for Pure Vision Based GUI Agent
|
||||
|
||||
<p align="center">
|
||||
<img src="imgs/logo.png" alt="Logo">
|
||||
</p>
|
||||
|
||||
[](https://arxiv.org/abs/2408.00203)
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
|
||||
📢 [[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)]
|
||||
|
||||
**OmniParser** is a comprehensive method for parsing user interface screenshots into structured and easy-to-understand elements, which significantly enhances the ability of GPT-4V to generate actions that can be accurately grounded in the corresponding regions of the interface.
|
||||
|
||||
## News
|
||||
- [2024/10] Both Interactive Region Detection Model and Icon functional description model are released! [Hugginface models](https://huggingface.co/microsoft/OmniParser)
|
||||
- [2024/09] OmniParser achieves the best performance on [Windows Agent Arena](https://microsoft.github.io/WindowsAgentArena/)!
|
||||
|
||||
## Install
|
||||
Install environment:
|
||||
```python
|
||||
conda create -n "omni" python==3.12
|
||||
conda activate omni
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
Then download the model ckpts files in: https://huggingface.co/microsoft/OmniParser, and put them under weights/, default folder structure is: weights/icon_detect, weights/icon_caption_florence, weights/icon_caption_blip2.
|
||||
|
||||
Finally, convert the safetensor to .pt file.
|
||||
```python
|
||||
python weights/convert_safetensor_to_pt.py
|
||||
```
|
||||
|
||||
## Examples:
|
||||
We put together a few simple examples in the demo.ipynb.
|
||||
|
||||
## Gradio Demo
|
||||
To run gradio demo, simply run:
|
||||
```python
|
||||
python gradio_demo.py
|
||||
```
|
||||
|
||||
|
||||
## 📚 Citation
|
||||
Our technical report can be found [here](https://arxiv.org/abs/2408.00203).
|
||||
If you find our work useful, please consider citing our work:
|
||||
```
|
||||
@misc{lu2024omniparserpurevisionbased,
|
||||
title={OmniParser for Pure Vision Based GUI Agent},
|
||||
author={Yadong Lu and Jianwei Yang and Yelong Shen and Ahmed Awadallah},
|
||||
year={2024},
|
||||
eprint={2408.00203},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CV},
|
||||
url={https://arxiv.org/abs/2408.00203},
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,174 @@
|
||||
---
|
||||
language: en
|
||||
license: mit
|
||||
tags:
|
||||
- vision
|
||||
- image-to-text
|
||||
- image-captioning
|
||||
- visual-question-answering
|
||||
pipeline_tag: image-to-text
|
||||
---
|
||||
|
||||
# BLIP-2, OPT-2.7b, pre-trained only
|
||||
|
||||
BLIP-2 model, leveraging [OPT-2.7b](https://huggingface.co/facebook/opt-2.7b) (a large language model with 2.7 billion parameters).
|
||||
It was introduced in the paper [BLIP-2: Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models](https://arxiv.org/abs/2301.12597) by Li et al. and first released in [this repository](https://github.com/salesforce/LAVIS/tree/main/projects/blip2).
|
||||
|
||||
Disclaimer: The team releasing BLIP-2 did not write a model card for this model so this model card has been written by the Hugging Face team.
|
||||
|
||||
## Model description
|
||||
|
||||
BLIP-2 consists of 3 models: a CLIP-like image encoder, a Querying Transformer (Q-Former) and a large language model.
|
||||
|
||||
The authors initialize the weights of the image encoder and large language model from pre-trained checkpoints and keep them frozen
|
||||
while training the Querying Transformer, which is a BERT-like Transformer encoder that maps a set of "query tokens" to query embeddings,
|
||||
which bridge the gap between the embedding space of the image encoder and the large language model.
|
||||
|
||||
The goal for the model is simply to predict the next text token, giving the query embeddings and the previous text.
|
||||
|
||||
<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/blip2_architecture.jpg"
|
||||
alt="drawing" width="600"/>
|
||||
|
||||
This allows the model to be used for tasks like:
|
||||
|
||||
- image captioning
|
||||
- visual question answering (VQA)
|
||||
- chat-like conversations by feeding the image and the previous conversation as prompt to the model
|
||||
|
||||
## Direct Use and Downstream Use
|
||||
|
||||
You can use the raw model for conditional text generation given an image and optional text. See the [model hub](https://huggingface.co/models?search=Salesforce/blip) to look for
|
||||
fine-tuned versions on a task that interests you.
|
||||
|
||||
## Bias, Risks, Limitations, and Ethical Considerations
|
||||
|
||||
BLIP2-OPT uses off-the-shelf OPT as the language model. It inherits the same risks and limitations as mentioned in Meta's model card.
|
||||
|
||||
> Like other large language models for which the diversity (or lack thereof) of training
|
||||
> data induces downstream impact on the quality of our model, OPT-175B has limitations in terms
|
||||
> of bias and safety. OPT-175B can also have quality issues in terms of generation diversity and
|
||||
> hallucination. In general, OPT-175B is not immune from the plethora of issues that plague modern
|
||||
> large language models.
|
||||
>
|
||||
BLIP2 is fine-tuned on image-text datasets (e.g. [LAION](https://laion.ai/blog/laion-400-open-dataset/) ) collected from the internet. As a result the model itself is potentially vulnerable to generating equivalently inappropriate content or replicating inherent biases in the underlying data.
|
||||
|
||||
BLIP2 has not been tested in real world applications. It should not be directly deployed in any applications. Researchers should first carefully assess the safety and fairness of the model in relation to the specific context they’re being deployed within.
|
||||
|
||||
|
||||
### How to use
|
||||
|
||||
For code examples, we refer to the [documentation](https://huggingface.co/docs/transformers/main/en/model_doc/blip-2#transformers.Blip2ForConditionalGeneration.forward.example).
|
||||
|
||||
### Memory requirements
|
||||
|
||||
The memory requirements differ based on the precision one uses. One can use 4-bit inference using [Bitsandbytes](https://huggingface.co/blog/4bit-transformers-bitsandbytes), which greatly reduce the memory requirements.
|
||||
|
||||
| dtype | Largest Layer or Residual Group | Total Size | Training using Adam |
|
||||
|-------------------|---------------------------------|------------|----------------------|
|
||||
| float32 | 490.94 MB | 14.43 GB | 57.72 GB |
|
||||
| float16/bfloat16 | 245.47 MB | 7.21 GB | 28.86 GB |
|
||||
| int8 | 122.73 MB | 3.61 GB | 14.43 GB |
|
||||
| int4 | 61.37 MB | 1.8 GB | 7.21 GB |
|
||||
|
||||
#### Running the model on CPU
|
||||
|
||||
<details>
|
||||
<summary> Click to expand </summary>
|
||||
|
||||
```python
|
||||
import requests
|
||||
from PIL import Image
|
||||
from transformers import Blip2Processor, Blip2ForConditionalGeneration
|
||||
|
||||
processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
|
||||
model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b")
|
||||
|
||||
img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg'
|
||||
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
|
||||
|
||||
question = "how many dogs are in the picture?"
|
||||
inputs = processor(raw_image, question, return_tensors="pt")
|
||||
|
||||
out = model.generate(**inputs)
|
||||
print(processor.decode(out[0], skip_special_tokens=True).strip())
|
||||
```
|
||||
</details>
|
||||
|
||||
#### Running the model on GPU
|
||||
|
||||
##### In full precision
|
||||
|
||||
<details>
|
||||
<summary> Click to expand </summary>
|
||||
|
||||
```python
|
||||
# pip install accelerate
|
||||
import requests
|
||||
from PIL import Image
|
||||
from transformers import Blip2Processor, Blip2ForConditionalGeneration
|
||||
|
||||
processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
|
||||
model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", device_map="auto")
|
||||
|
||||
img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg'
|
||||
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
|
||||
|
||||
question = "how many dogs are in the picture?"
|
||||
inputs = processor(raw_image, question, return_tensors="pt").to("cuda")
|
||||
|
||||
out = model.generate(**inputs)
|
||||
print(processor.decode(out[0], skip_special_tokens=True).strip())
|
||||
```
|
||||
</details>
|
||||
|
||||
##### In half precision (`float16`)
|
||||
|
||||
<details>
|
||||
<summary> Click to expand </summary>
|
||||
|
||||
```python
|
||||
# pip install accelerate
|
||||
import torch
|
||||
import requests
|
||||
from PIL import Image
|
||||
from transformers import Blip2Processor, Blip2ForConditionalGeneration
|
||||
|
||||
processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
|
||||
model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16, device_map="auto")
|
||||
|
||||
img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg'
|
||||
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
|
||||
|
||||
question = "how many dogs are in the picture?"
|
||||
inputs = processor(raw_image, question, return_tensors="pt").to("cuda", torch.float16)
|
||||
|
||||
out = model.generate(**inputs)
|
||||
print(processor.decode(out[0], skip_special_tokens=True).strip())
|
||||
```
|
||||
</details>
|
||||
|
||||
##### In 8-bit precision (`int8`)
|
||||
|
||||
<details>
|
||||
<summary> Click to expand </summary>
|
||||
|
||||
```python
|
||||
# pip install accelerate bitsandbytes
|
||||
import torch
|
||||
import requests
|
||||
from PIL import Image
|
||||
from transformers import Blip2Processor, Blip2ForConditionalGeneration
|
||||
|
||||
processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
|
||||
model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", load_in_8bit=True, device_map="auto")
|
||||
|
||||
img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg'
|
||||
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
|
||||
|
||||
question = "how many dogs are in the picture?"
|
||||
inputs = processor(raw_image, question, return_tensors="pt").to("cuda", torch.float16)
|
||||
|
||||
out = model.generate(**inputs)
|
||||
print(processor.decode(out[0], skip_special_tokens=True).strip())
|
||||
```
|
||||
</details>
|
||||
@@ -0,0 +1,255 @@
|
||||
{
|
||||
"_commit_hash": null,
|
||||
"architectures": [
|
||||
"Blip2ForConditionalGeneration"
|
||||
],
|
||||
"initializer_factor": 1.0,
|
||||
"initializer_range": 0.02,
|
||||
"model_type": "blip-2",
|
||||
"num_query_tokens": 32,
|
||||
"qformer_config": {
|
||||
"_name_or_path": "",
|
||||
"add_cross_attention": false,
|
||||
"architectures": null,
|
||||
"attention_probs_dropout_prob": 0.1,
|
||||
"bad_words_ids": null,
|
||||
"begin_suppress_tokens": null,
|
||||
"bos_token_id": null,
|
||||
"chunk_size_feed_forward": 0,
|
||||
"classifier_dropout": null,
|
||||
"cross_attention_frequency": 2,
|
||||
"cross_attention_hidden_size": null,
|
||||
"decoder_start_token_id": null,
|
||||
"diversity_penalty": 0.0,
|
||||
"do_sample": false,
|
||||
"early_stopping": false,
|
||||
"encoder_hidden_size": 1408,
|
||||
"encoder_no_repeat_ngram_size": 0,
|
||||
"eos_token_id": null,
|
||||
"exponential_decay_length_penalty": null,
|
||||
"finetuning_task": null,
|
||||
"forced_bos_token_id": null,
|
||||
"forced_eos_token_id": null,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_dropout_prob": 0.1,
|
||||
"hidden_size": 768,
|
||||
"id2label": {
|
||||
"0": "LABEL_0",
|
||||
"1": "LABEL_1"
|
||||
},
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 3072,
|
||||
"is_decoder": false,
|
||||
"is_encoder_decoder": false,
|
||||
"label2id": {
|
||||
"LABEL_0": 0,
|
||||
"LABEL_1": 1
|
||||
},
|
||||
"layer_norm_eps": 1e-12,
|
||||
"length_penalty": 1.0,
|
||||
"max_length": 20,
|
||||
"max_position_embeddings": 512,
|
||||
"min_length": 0,
|
||||
"model_type": "blip_2_qformer",
|
||||
"no_repeat_ngram_size": 0,
|
||||
"num_attention_heads": 12,
|
||||
"num_beam_groups": 1,
|
||||
"num_beams": 1,
|
||||
"num_hidden_layers": 12,
|
||||
"num_return_sequences": 1,
|
||||
"output_attentions": false,
|
||||
"output_hidden_states": false,
|
||||
"output_scores": false,
|
||||
"pad_token_id": 0,
|
||||
"position_embedding_type": "absolute",
|
||||
"prefix": null,
|
||||
"problem_type": null,
|
||||
"pruned_heads": {},
|
||||
"remove_invalid_values": false,
|
||||
"repetition_penalty": 1.0,
|
||||
"return_dict": true,
|
||||
"return_dict_in_generate": false,
|
||||
"sep_token_id": null,
|
||||
"suppress_tokens": null,
|
||||
"task_specific_params": null,
|
||||
"temperature": 1.0,
|
||||
"tf_legacy_loss": false,
|
||||
"tie_encoder_decoder": false,
|
||||
"tie_word_embeddings": true,
|
||||
"tokenizer_class": null,
|
||||
"top_k": 50,
|
||||
"top_p": 1.0,
|
||||
"torch_dtype": null,
|
||||
"torchscript": false,
|
||||
"transformers_version": "4.27.0.dev0",
|
||||
"typical_p": 1.0,
|
||||
"use_bfloat16": false,
|
||||
"vocab_size": 30522
|
||||
},
|
||||
"text_config": {
|
||||
"_name_or_path": "facebook/opt-2.7b",
|
||||
"_remove_final_layer_norm": false,
|
||||
"activation_dropout": 0.0,
|
||||
"activation_function": "relu",
|
||||
"add_cross_attention": false,
|
||||
"architectures": [
|
||||
"OPTForCausalLM"
|
||||
],
|
||||
"attention_dropout": 0.0,
|
||||
"bad_words_ids": null,
|
||||
"begin_suppress_tokens": null,
|
||||
"bos_token_id": 2,
|
||||
"chunk_size_feed_forward": 0,
|
||||
"cross_attention_hidden_size": null,
|
||||
"decoder_start_token_id": null,
|
||||
"diversity_penalty": 0.0,
|
||||
"do_layer_norm_before": true,
|
||||
"do_sample": false,
|
||||
"dropout": 0.1,
|
||||
"early_stopping": false,
|
||||
"enable_bias": true,
|
||||
"encoder_no_repeat_ngram_size": 0,
|
||||
"eos_token_id": 50118,
|
||||
"exponential_decay_length_penalty": null,
|
||||
"ffn_dim": 10240,
|
||||
"finetuning_task": null,
|
||||
"forced_bos_token_id": null,
|
||||
"forced_eos_token_id": null,
|
||||
"hidden_size": 2560,
|
||||
"id2label": {
|
||||
"0": "LABEL_0",
|
||||
"1": "LABEL_1"
|
||||
},
|
||||
"init_std": 0.02,
|
||||
"is_decoder": false,
|
||||
"is_encoder_decoder": false,
|
||||
"label2id": {
|
||||
"LABEL_0": 0,
|
||||
"LABEL_1": 1
|
||||
},
|
||||
"layer_norm_elementwise_affine": true,
|
||||
"layerdrop": 0.0,
|
||||
"length_penalty": 1.0,
|
||||
"max_length": 20,
|
||||
"max_position_embeddings": 2048,
|
||||
"min_length": 0,
|
||||
"model_type": "opt",
|
||||
"no_repeat_ngram_size": 0,
|
||||
"num_attention_heads": 32,
|
||||
"num_beam_groups": 1,
|
||||
"num_beams": 1,
|
||||
"num_hidden_layers": 32,
|
||||
"num_return_sequences": 1,
|
||||
"output_attentions": false,
|
||||
"output_hidden_states": false,
|
||||
"output_scores": false,
|
||||
"pad_token_id": 1,
|
||||
"prefix": "</s>",
|
||||
"problem_type": null,
|
||||
"pruned_heads": {},
|
||||
"remove_invalid_values": false,
|
||||
"repetition_penalty": 1.0,
|
||||
"return_dict": true,
|
||||
"return_dict_in_generate": false,
|
||||
"sep_token_id": null,
|
||||
"suppress_tokens": null,
|
||||
"task_specific_params": null,
|
||||
"temperature": 1.0,
|
||||
"tf_legacy_loss": false,
|
||||
"tie_encoder_decoder": false,
|
||||
"tie_word_embeddings": true,
|
||||
"tokenizer_class": null,
|
||||
"top_k": 50,
|
||||
"top_p": 1.0,
|
||||
"torch_dtype": "float16",
|
||||
"torchscript": false,
|
||||
"transformers_version": "4.27.0.dev0",
|
||||
"typical_p": 1.0,
|
||||
"use_bfloat16": false,
|
||||
"use_cache": true,
|
||||
"vocab_size": 50272,
|
||||
"word_embed_proj_dim": 2560
|
||||
},
|
||||
"torch_dtype": "float32",
|
||||
"transformers_version": null,
|
||||
"use_decoder_only_language_model": true,
|
||||
"vision_config": {
|
||||
"_name_or_path": "",
|
||||
"add_cross_attention": false,
|
||||
"architectures": null,
|
||||
"attention_dropout": 0.0,
|
||||
"bad_words_ids": null,
|
||||
"begin_suppress_tokens": null,
|
||||
"bos_token_id": null,
|
||||
"chunk_size_feed_forward": 0,
|
||||
"cross_attention_hidden_size": null,
|
||||
"decoder_start_token_id": null,
|
||||
"diversity_penalty": 0.0,
|
||||
"do_sample": false,
|
||||
"dropout": 0.0,
|
||||
"early_stopping": false,
|
||||
"encoder_no_repeat_ngram_size": 0,
|
||||
"eos_token_id": null,
|
||||
"exponential_decay_length_penalty": null,
|
||||
"finetuning_task": null,
|
||||
"forced_bos_token_id": null,
|
||||
"forced_eos_token_id": null,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_size": 1408,
|
||||
"id2label": {
|
||||
"0": "LABEL_0",
|
||||
"1": "LABEL_1"
|
||||
},
|
||||
"image_size": 224,
|
||||
"initializer_factor": 1.0,
|
||||
"initializer_range": 1e-10,
|
||||
"intermediate_size": 6144,
|
||||
"is_decoder": false,
|
||||
"is_encoder_decoder": false,
|
||||
"label2id": {
|
||||
"LABEL_0": 0,
|
||||
"LABEL_1": 1
|
||||
},
|
||||
"layer_norm_eps": 1e-6,
|
||||
"length_penalty": 1.0,
|
||||
"max_length": 20,
|
||||
"min_length": 0,
|
||||
"model_type": "blip_2_vision_model",
|
||||
"no_repeat_ngram_size": 0,
|
||||
"num_attention_heads": 16,
|
||||
"num_beam_groups": 1,
|
||||
"num_beams": 1,
|
||||
"num_channels": 3,
|
||||
"num_hidden_layers": 39,
|
||||
"num_return_sequences": 1,
|
||||
"output_attentions": false,
|
||||
"output_hidden_states": false,
|
||||
"output_scores": false,
|
||||
"pad_token_id": null,
|
||||
"patch_size": 14,
|
||||
"prefix": null,
|
||||
"problem_type": null,
|
||||
"projection_dim": 512,
|
||||
"pruned_heads": {},
|
||||
"qkv_bias": true,
|
||||
"remove_invalid_values": false,
|
||||
"repetition_penalty": 1.0,
|
||||
"return_dict": true,
|
||||
"return_dict_in_generate": false,
|
||||
"sep_token_id": null,
|
||||
"suppress_tokens": null,
|
||||
"task_specific_params": null,
|
||||
"temperature": 1.0,
|
||||
"tf_legacy_loss": false,
|
||||
"tie_encoder_decoder": false,
|
||||
"tie_word_embeddings": true,
|
||||
"tokenizer_class": null,
|
||||
"top_k": 50,
|
||||
"top_p": 1.0,
|
||||
"torch_dtype": null,
|
||||
"torchscript": false,
|
||||
"transformers_version": "4.27.0.dev0",
|
||||
"typical_p": 1.0,
|
||||
"use_bfloat16": false
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"do_convert_rgb": true,
|
||||
"do_normalize": true,
|
||||
"do_rescale": true,
|
||||
"do_resize": true,
|
||||
"image_mean": [
|
||||
0.48145466,
|
||||
0.4578275,
|
||||
0.40821073
|
||||
],
|
||||
"image_processor_type": "BlipImageProcessor",
|
||||
"image_std": [
|
||||
0.26862954,
|
||||
0.26130258,
|
||||
0.27577711
|
||||
],
|
||||
"processor_class": "Blip2Processor",
|
||||
"resample": 3,
|
||||
"rescale_factor": 0.00392156862745098,
|
||||
"size": {
|
||||
"height": 224,
|
||||
"width": 224
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,30 @@
|
||||
{
|
||||
"bos_token": {
|
||||
"content": "</s>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"eos_token": {
|
||||
"content": "</s>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"pad_token": {
|
||||
"content": "<pad>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"unk_token": {
|
||||
"content": "</s>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,41 @@
|
||||
{
|
||||
"add_bos_token": true,
|
||||
"add_prefix_space": false,
|
||||
"bos_token": {
|
||||
"__type": "AddedToken",
|
||||
"content": "</s>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"eos_token": {
|
||||
"__type": "AddedToken",
|
||||
"content": "</s>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"errors": "replace",
|
||||
"model_max_length": 1000000000000000019884624838656,
|
||||
"pad_token": {
|
||||
"__type": "AddedToken",
|
||||
"content": "<pad>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"processor_class": "Blip2Processor",
|
||||
"special_tokens_map_file": null,
|
||||
"tokenizer_class": "GPT2Tokenizer",
|
||||
"unk_token": {
|
||||
"__type": "AddedToken",
|
||||
"content": "</s>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
}
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,60 @@
|
||||
from utils import get_som_labeled_img, check_ocr_box, get_caption_model_processor, get_dino_model, get_yolo_model
|
||||
import torch
|
||||
from ultralytics import YOLO
|
||||
from PIL import Image
|
||||
from typing import Dict, Tuple, List
|
||||
import io
|
||||
import base64
|
||||
|
||||
|
||||
config = {
|
||||
'som_model_path': 'finetuned_icon_detect.pt',
|
||||
'device': 'cpu',
|
||||
'caption_model_path': 'Salesforce/blip2-opt-2.7b',
|
||||
'draw_bbox_config': {
|
||||
'text_scale': 0.8,
|
||||
'text_thickness': 2,
|
||||
'text_padding': 3,
|
||||
'thickness': 3,
|
||||
},
|
||||
'BOX_TRESHOLD': 0.05
|
||||
}
|
||||
|
||||
|
||||
class Omniparser(object):
|
||||
def __init__(self, config: Dict):
|
||||
self.config = config
|
||||
|
||||
self.som_model = get_yolo_model(model_path=config['som_model_path'])
|
||||
# self.caption_model_processor = get_caption_model_processor(config['caption_model_path'], device=cofig['device'])
|
||||
# self.caption_model_processor['model'].to(torch.float32)
|
||||
|
||||
def parse(self, image_path: str):
|
||||
print('Parsing image:', image_path)
|
||||
ocr_bbox_rslt, is_goal_filtered = check_ocr_box(image_path, display_img = False, output_bb_format='xyxy', goal_filtering=None, easyocr_args={'paragraph': False, 'text_threshold':0.9})
|
||||
text, ocr_bbox = ocr_bbox_rslt
|
||||
|
||||
draw_bbox_config = self.config['draw_bbox_config']
|
||||
BOX_TRESHOLD = self.config['BOX_TRESHOLD']
|
||||
dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(image_path, self.som_model, BOX_TRESHOLD = BOX_TRESHOLD, output_coord_in_ratio=False, ocr_bbox=ocr_bbox,draw_bbox_config=draw_bbox_config, caption_model_processor=None, ocr_text=text,use_local_semantics=False)
|
||||
|
||||
image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))
|
||||
# formating output
|
||||
return_list = [{'from': 'omniparser', 'shape': {'x':coord[0], 'y':coord[1], 'width':coord[2], 'height':coord[3]},
|
||||
'text': parsed_content_list[i].split(': ')[1], 'type':'text'} for i, (k, coord) in enumerate(label_coordinates.items()) if i < len(parsed_content_list)]
|
||||
return_list.extend(
|
||||
[{'from': 'omniparser', 'shape': {'x':coord[0], 'y':coord[1], 'width':coord[2], 'height':coord[3]},
|
||||
'text': 'None', 'type':'icon'} for i, (k, coord) in enumerate(label_coordinates.items()) if i >= len(parsed_content_list)]
|
||||
)
|
||||
|
||||
return [image, return_list]
|
||||
|
||||
parser = Omniparser(config)
|
||||
image_path = 'examples/pc_1.png'
|
||||
|
||||
# time the parser
|
||||
import time
|
||||
s = time.time()
|
||||
image, parsed_content_list = parser.parse(image_path)
|
||||
device = config['device']
|
||||
print(f'Time taken for Omniparser on {device}:', time.time() - s)
|
||||
@@ -0,0 +1,425 @@
|
||||
'''
|
||||
Adapted from https://github.com/google-research/google-research/tree/master/android_in_the_wild
|
||||
'''
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import numpy as np
|
||||
|
||||
# import action_type as action_type_lib
|
||||
import enum
|
||||
|
||||
class ActionType(enum.IntEnum):
|
||||
# Placeholders for unused enum values
|
||||
UNUSED_0 = 0
|
||||
UNUSED_1 = 1
|
||||
UNUSED_2 = 2
|
||||
UNUSED_8 = 8
|
||||
UNUSED_9 = 9
|
||||
|
||||
########### Agent actions ###########
|
||||
|
||||
# A type action that sends text to the emulator. Note that this simply sends
|
||||
# text and does not perform any clicks for element focus or enter presses for
|
||||
# submitting text.
|
||||
TYPE = 3
|
||||
|
||||
# The dual point action used to represent all gestures.
|
||||
DUAL_POINT = 4
|
||||
|
||||
# These actions differentiate pressing the home and back button from touches.
|
||||
# They represent explicit presses of back and home performed using ADB.
|
||||
PRESS_BACK = 5
|
||||
PRESS_HOME = 6
|
||||
|
||||
# An action representing that ADB command for hitting enter was performed.
|
||||
PRESS_ENTER = 7
|
||||
|
||||
########### Episode status actions ###########
|
||||
|
||||
# An action used to indicate the desired task has been completed and resets
|
||||
# the environment. This action should also be used in the case that the task
|
||||
# has already been completed and there is nothing to do.
|
||||
# e.g. The task is to turn on the Wi-Fi when it is already on
|
||||
STATUS_TASK_COMPLETE = 10
|
||||
|
||||
# An action used to indicate that desired task is impossible to complete and
|
||||
# resets the environment. This can be a result of many different things
|
||||
# including UI changes, Android version differences, etc.
|
||||
STATUS_TASK_IMPOSSIBLE = 11
|
||||
|
||||
|
||||
_TAP_DISTANCE_THRESHOLD = 0.14 # Fraction of the screen
|
||||
ANNOTATION_WIDTH_AUGMENT_FRACTION = 1.4
|
||||
ANNOTATION_HEIGHT_AUGMENT_FRACTION = 1.4
|
||||
|
||||
# Interval determining if an action is a tap or a swipe.
|
||||
_SWIPE_DISTANCE_THRESHOLD = 0.04
|
||||
|
||||
|
||||
def _yx_in_bounding_boxes(
|
||||
yx, bounding_boxes
|
||||
):
|
||||
"""Check if the (y,x) point is contained in each bounding box.
|
||||
|
||||
Args:
|
||||
yx: The (y, x) coordinate in pixels of the point.
|
||||
bounding_boxes: A 2D int array of shape (num_bboxes, 4), where each row
|
||||
represents a bounding box: (y_top_left, x_top_left, box_height,
|
||||
box_width). Note: containment is inclusive of the bounding box edges.
|
||||
|
||||
Returns:
|
||||
is_inside: A 1D bool array where each element specifies if the point is
|
||||
contained within the respective box.
|
||||
"""
|
||||
y, x = yx
|
||||
|
||||
# `bounding_boxes` has shape (n_elements, 4); we extract each array along the
|
||||
# last axis into shape (n_elements, 1), then squeeze unneeded dimension.
|
||||
top, left, height, width = [
|
||||
jnp.squeeze(v, axis=-1) for v in jnp.split(bounding_boxes, 4, axis=-1)
|
||||
]
|
||||
|
||||
# The y-axis is inverted for AndroidEnv, so bottom = top + height.
|
||||
bottom, right = top + height, left + width
|
||||
|
||||
return jnp.logical_and(y >= top, y <= bottom) & jnp.logical_and(
|
||||
x >= left, x <= right)
|
||||
|
||||
|
||||
def _resize_annotation_bounding_boxes(
|
||||
annotation_positions, annotation_width_augment_fraction,
|
||||
annotation_height_augment_fraction):
|
||||
"""Resize the bounding boxes by the given fractions.
|
||||
|
||||
Args:
|
||||
annotation_positions: Array of shape (N, 4), where each row represents the
|
||||
(y, x, height, width) of the bounding boxes.
|
||||
annotation_width_augment_fraction: The fraction to augment the box widths,
|
||||
E.g., 1.4 == 240% total increase.
|
||||
annotation_height_augment_fraction: Same as described for width, but for box
|
||||
height.
|
||||
|
||||
Returns:
|
||||
Resized bounding box.
|
||||
|
||||
"""
|
||||
height_change = (
|
||||
annotation_height_augment_fraction * annotation_positions[:, 2])
|
||||
width_change = (
|
||||
annotation_width_augment_fraction * annotation_positions[:, 3])
|
||||
|
||||
# Limit bounding box positions to the screen.
|
||||
resized_annotations = jnp.stack([
|
||||
jnp.maximum(0, annotation_positions[:, 0] - (height_change / 2)),
|
||||
jnp.maximum(0, annotation_positions[:, 1] - (width_change / 2)),
|
||||
jnp.minimum(1, annotation_positions[:, 2] + height_change),
|
||||
jnp.minimum(1, annotation_positions[:, 3] + width_change),
|
||||
],
|
||||
axis=1)
|
||||
return resized_annotations
|
||||
|
||||
|
||||
def is_tap_action(normalized_start_yx,
|
||||
normalized_end_yx):
|
||||
distance = jnp.linalg.norm(
|
||||
jnp.array(normalized_start_yx) - jnp.array(normalized_end_yx))
|
||||
return distance <= _SWIPE_DISTANCE_THRESHOLD
|
||||
|
||||
|
||||
def _is_non_dual_point_action(action_type):
|
||||
return jnp.not_equal(action_type, ActionType.DUAL_POINT)
|
||||
|
||||
|
||||
def _check_tap_actions_match(
|
||||
tap_1_yx,
|
||||
tap_2_yx,
|
||||
annotation_positions,
|
||||
matching_tap_distance_threshold_screen_percentage,
|
||||
annotation_width_augment_fraction,
|
||||
annotation_height_augment_fraction,
|
||||
):
|
||||
"""Determines if two tap actions are the same."""
|
||||
resized_annotation_positions = _resize_annotation_bounding_boxes(
|
||||
annotation_positions,
|
||||
annotation_width_augment_fraction,
|
||||
annotation_height_augment_fraction,
|
||||
)
|
||||
|
||||
# Check if the ground truth tap action falls in an annotation's bounding box.
|
||||
tap1_in_box = _yx_in_bounding_boxes(tap_1_yx, resized_annotation_positions)
|
||||
tap2_in_box = _yx_in_bounding_boxes(tap_2_yx, resized_annotation_positions)
|
||||
both_in_box = jnp.max(tap1_in_box & tap2_in_box)
|
||||
|
||||
# If the ground-truth tap action falls outside any of the annotation
|
||||
# bounding boxes or one of the actions is inside a bounding box and the other
|
||||
# is outside bounding box or vice versa, compare the points using Euclidean
|
||||
# distance.
|
||||
within_threshold = (
|
||||
jnp.linalg.norm(jnp.array(tap_1_yx) - jnp.array(tap_2_yx))
|
||||
<= matching_tap_distance_threshold_screen_percentage
|
||||
)
|
||||
return jnp.logical_or(both_in_box, within_threshold)
|
||||
|
||||
|
||||
def _check_drag_actions_match(
|
||||
drag_1_touch_yx,
|
||||
drag_1_lift_yx,
|
||||
drag_2_touch_yx,
|
||||
drag_2_lift_yx,
|
||||
):
|
||||
"""Determines if two drag actions are the same."""
|
||||
# Store drag deltas (the change in the y and x coordinates from touch to
|
||||
# lift), magnitudes, and the index of the main axis, which is the axis with
|
||||
# the greatest change in coordinate value (e.g. a drag starting at (0, 0) and
|
||||
# ending at (0.3, 0.5) has a main axis index of 1).
|
||||
drag_1_deltas = drag_1_lift_yx - drag_1_touch_yx
|
||||
drag_1_magnitudes = jnp.abs(drag_1_deltas)
|
||||
drag_1_main_axis = np.argmax(drag_1_magnitudes)
|
||||
drag_2_deltas = drag_2_lift_yx - drag_2_touch_yx
|
||||
drag_2_magnitudes = jnp.abs(drag_2_deltas)
|
||||
drag_2_main_axis = np.argmax(drag_2_magnitudes)
|
||||
|
||||
return jnp.equal(drag_1_main_axis, drag_2_main_axis)
|
||||
|
||||
|
||||
def check_actions_match(
|
||||
action_1_touch_yx,
|
||||
action_1_lift_yx,
|
||||
action_1_action_type,
|
||||
action_2_touch_yx,
|
||||
action_2_lift_yx,
|
||||
action_2_action_type,
|
||||
annotation_positions,
|
||||
tap_distance_threshold = _TAP_DISTANCE_THRESHOLD,
|
||||
annotation_width_augment_fraction = ANNOTATION_WIDTH_AUGMENT_FRACTION,
|
||||
annotation_height_augment_fraction = ANNOTATION_HEIGHT_AUGMENT_FRACTION,
|
||||
):
|
||||
"""Determines if two actions are considered to be the same.
|
||||
|
||||
Two actions being "the same" is defined here as two actions that would result
|
||||
in a similar screen state.
|
||||
|
||||
Args:
|
||||
action_1_touch_yx: The (y, x) coordinates of the first action's touch.
|
||||
action_1_lift_yx: The (y, x) coordinates of the first action's lift.
|
||||
action_1_action_type: The action type of the first action.
|
||||
action_2_touch_yx: The (y, x) coordinates of the second action's touch.
|
||||
action_2_lift_yx: The (y, x) coordinates of the second action's lift.
|
||||
action_2_action_type: The action type of the second action.
|
||||
annotation_positions: The positions of the UI annotations for the screen. It
|
||||
is A 2D int array of shape (num_bboxes, 4), where each row represents a
|
||||
bounding box: (y_top_left, x_top_left, box_height, box_width). Note that
|
||||
containment is inclusive of the bounding box edges.
|
||||
tap_distance_threshold: The threshold that determines if two taps result in
|
||||
a matching screen state if they don't fall the same bounding boxes.
|
||||
annotation_width_augment_fraction: The fraction to increase the width of the
|
||||
bounding box by.
|
||||
annotation_height_augment_fraction: The fraction to increase the height of
|
||||
of the bounding box by.
|
||||
|
||||
Returns:
|
||||
A boolean representing whether the two given actions are the same or not.
|
||||
"""
|
||||
action_1_touch_yx = jnp.asarray(action_1_touch_yx)
|
||||
action_1_lift_yx = jnp.asarray(action_1_lift_yx)
|
||||
action_2_touch_yx = jnp.asarray(action_2_touch_yx)
|
||||
action_2_lift_yx = jnp.asarray(action_2_lift_yx)
|
||||
|
||||
# Checks if at least one of the actions is global (i.e. not DUAL_POINT),
|
||||
# because if that is the case, only the actions' types need to be compared.
|
||||
has_non_dual_point_action = jnp.logical_or(
|
||||
_is_non_dual_point_action(action_1_action_type),
|
||||
_is_non_dual_point_action(action_2_action_type),
|
||||
)
|
||||
#print("non dual point: "+str(has_non_dual_point_action))
|
||||
|
||||
different_dual_point_types = jnp.logical_xor(
|
||||
is_tap_action(action_1_touch_yx, action_1_lift_yx),
|
||||
is_tap_action(action_2_touch_yx, action_2_lift_yx),
|
||||
)
|
||||
#print("different dual type: "+str(different_dual_point_types))
|
||||
|
||||
is_tap = jnp.logical_and(
|
||||
is_tap_action(action_1_touch_yx, action_1_lift_yx),
|
||||
is_tap_action(action_2_touch_yx, action_2_lift_yx),
|
||||
)
|
||||
#print("is tap: "+str(is_tap))
|
||||
|
||||
taps_match = _check_tap_actions_match(
|
||||
action_1_touch_yx,
|
||||
action_2_touch_yx,
|
||||
annotation_positions,
|
||||
tap_distance_threshold,
|
||||
annotation_width_augment_fraction,
|
||||
annotation_height_augment_fraction,
|
||||
)
|
||||
#print("tap match: "+str(taps_match))
|
||||
|
||||
taps_match = jnp.logical_and(is_tap, taps_match)
|
||||
#print("tap match: "+str(taps_match))
|
||||
|
||||
drags_match = _check_drag_actions_match(
|
||||
action_1_touch_yx, action_1_lift_yx, action_2_touch_yx, action_2_lift_yx
|
||||
)
|
||||
drags_match = jnp.where(is_tap, False, drags_match)
|
||||
#print("drag match: "+str(drags_match))
|
||||
|
||||
return jnp.where(
|
||||
has_non_dual_point_action,
|
||||
jnp.equal(action_1_action_type, action_2_action_type),
|
||||
jnp.where(
|
||||
different_dual_point_types,
|
||||
False,
|
||||
jnp.logical_or(taps_match, drags_match),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def action_2_format(step_data):
|
||||
# 把test数据集中的动作格式转换为计算matching score的格式
|
||||
action_type = step_data["action_type_id"]
|
||||
|
||||
if action_type == 4:
|
||||
if step_data["action_type_text"] == 'click': # 点击
|
||||
touch_point = step_data["touch"]
|
||||
lift_point = step_data["lift"]
|
||||
else: # 上下左右滑动
|
||||
if step_data["action_type_text"] == 'scroll down':
|
||||
touch_point = [0.5, 0.8]
|
||||
lift_point = [0.5, 0.2]
|
||||
elif step_data["action_type_text"] == 'scroll up':
|
||||
touch_point = [0.5, 0.2]
|
||||
lift_point = [0.5, 0.8]
|
||||
elif step_data["action_type_text"] == 'scroll left':
|
||||
touch_point = [0.2, 0.5]
|
||||
lift_point = [0.8, 0.5]
|
||||
elif step_data["action_type_text"] == 'scroll right':
|
||||
touch_point = [0.8, 0.5]
|
||||
lift_point = [0.2, 0.5]
|
||||
else:
|
||||
touch_point = [-1.0, -1.0]
|
||||
lift_point = [-1.0, -1.0]
|
||||
|
||||
if action_type == 3:
|
||||
typed_text = step_data["type_text"]
|
||||
else:
|
||||
typed_text = ""
|
||||
|
||||
action = {"action_type": action_type, "touch_point": touch_point, "lift_point": lift_point,
|
||||
"typed_text": typed_text}
|
||||
|
||||
action["touch_point"] = [action["touch_point"][1], action["touch_point"][0]]
|
||||
action["lift_point"] = [action["lift_point"][1], action["lift_point"][0]]
|
||||
action["typed_text"] = action["typed_text"].lower()
|
||||
|
||||
return action
|
||||
|
||||
|
||||
def pred_2_format(step_data):
|
||||
# 把模型输出的内容转换为计算action_matching的格式
|
||||
action_type = step_data["action_type"]
|
||||
|
||||
if action_type == 4: # 点击
|
||||
action_type_new = 4
|
||||
touch_point = step_data["click_point"]
|
||||
lift_point = step_data["click_point"]
|
||||
typed_text = ""
|
||||
elif action_type == 0:
|
||||
action_type_new = 4
|
||||
touch_point = [0.5, 0.8]
|
||||
lift_point = [0.5, 0.2]
|
||||
typed_text = ""
|
||||
elif action_type == 1:
|
||||
action_type_new = 4
|
||||
touch_point = [0.5, 0.2]
|
||||
lift_point = [0.5, 0.8]
|
||||
typed_text = ""
|
||||
elif action_type == 8:
|
||||
action_type_new = 4
|
||||
touch_point = [0.2, 0.5]
|
||||
lift_point = [0.8, 0.5]
|
||||
typed_text = ""
|
||||
elif action_type == 9:
|
||||
action_type_new = 4
|
||||
touch_point = [0.8, 0.5]
|
||||
lift_point = [0.2, 0.5]
|
||||
typed_text = ""
|
||||
else:
|
||||
action_type_new = action_type
|
||||
touch_point = [-1.0, -1.0]
|
||||
lift_point = [-1.0, -1.0]
|
||||
typed_text = ""
|
||||
if action_type_new == 3:
|
||||
typed_text = step_data["typed_text"]
|
||||
|
||||
action = {"action_type": action_type_new, "touch_point": touch_point, "lift_point": lift_point,
|
||||
"typed_text": typed_text}
|
||||
|
||||
action["touch_point"] = [action["touch_point"][1], action["touch_point"][0]]
|
||||
action["lift_point"] = [action["lift_point"][1], action["lift_point"][0]]
|
||||
action["typed_text"] = action["typed_text"].lower()
|
||||
|
||||
return action
|
||||
|
||||
|
||||
def pred_2_format_simplified(step_data):
|
||||
# 把模型输出的内容转换为计算action_matching的格式
|
||||
action_type = step_data["action_type"]
|
||||
|
||||
if action_type == 'click' : # 点击
|
||||
action_type_new = 4
|
||||
touch_point = step_data["click_point"]
|
||||
lift_point = step_data["click_point"]
|
||||
typed_text = ""
|
||||
elif action_type == 'scroll' and step_data["direction"] == 'down':
|
||||
action_type_new = 4
|
||||
touch_point = [0.5, 0.8]
|
||||
lift_point = [0.5, 0.2]
|
||||
typed_text = ""
|
||||
elif action_type == 'scroll' and step_data["direction"] == 'up':
|
||||
action_type_new = 4
|
||||
touch_point = [0.5, 0.2]
|
||||
lift_point = [0.5, 0.8]
|
||||
typed_text = ""
|
||||
elif action_type == 'scroll' and step_data["direction"] == 'left':
|
||||
action_type_new = 4
|
||||
touch_point = [0.2, 0.5]
|
||||
lift_point = [0.8, 0.5]
|
||||
typed_text = ""
|
||||
elif action_type == 'scroll' and step_data["direction"] == 'right':
|
||||
action_type_new = 4
|
||||
touch_point = [0.8, 0.5]
|
||||
lift_point = [0.2, 0.5]
|
||||
typed_text = ""
|
||||
elif action_type == 'type':
|
||||
action_type_new = 3
|
||||
touch_point = [-1.0, -1.0]
|
||||
lift_point = [-1.0, -1.0]
|
||||
typed_text = step_data["text"]
|
||||
elif action_type == 'navigate_back':
|
||||
action_type_new = 5
|
||||
touch_point = [-1.0, -1.0]
|
||||
lift_point = [-1.0, -1.0]
|
||||
typed_text = ""
|
||||
elif action_type == 'navigate_home':
|
||||
action_type_new = 6
|
||||
touch_point = [-1.0, -1.0]
|
||||
lift_point = [-1.0, -1.0]
|
||||
typed_text = ""
|
||||
else:
|
||||
action_type_new = action_type
|
||||
touch_point = [-1.0, -1.0]
|
||||
lift_point = [-1.0, -1.0]
|
||||
typed_text = ""
|
||||
# if action_type_new == 'type':
|
||||
# typed_text = step_data["text"]
|
||||
|
||||
action = {"action_type": action_type_new, "touch_point": touch_point, "lift_point": lift_point,
|
||||
"typed_text": typed_text}
|
||||
|
||||
action["touch_point"] = [action["touch_point"][1], action["touch_point"][0]]
|
||||
action["lift_point"] = [action["lift_point"][1], action["lift_point"][0]]
|
||||
action["typed_text"] = action["typed_text"].lower()
|
||||
|
||||
return action
|
||||
@@ -0,0 +1,45 @@
|
||||
'''
|
||||
Adapted from https://github.com/google-research/google-research/tree/master/android_in_the_wild
|
||||
'''
|
||||
|
||||
import enum
|
||||
|
||||
class ActionType(enum.IntEnum):
|
||||
|
||||
# Placeholders for unused enum values
|
||||
UNUSED_0 = 0
|
||||
UNUSED_1 = 1
|
||||
UNUSED_2 = 2
|
||||
UNUSED_8 = 8
|
||||
UNUSED_9 = 9
|
||||
|
||||
########### Agent actions ###########
|
||||
|
||||
# A type action that sends text to the emulator. Note that this simply sends
|
||||
# text and does not perform any clicks for element focus or enter presses for
|
||||
# submitting text.
|
||||
TYPE = 3
|
||||
|
||||
# The dual point action used to represent all gestures.
|
||||
DUAL_POINT = 4
|
||||
|
||||
# These actions differentiate pressing the home and back button from touches.
|
||||
# They represent explicit presses of back and home performed using ADB.
|
||||
PRESS_BACK = 5
|
||||
PRESS_HOME = 6
|
||||
|
||||
# An action representing that ADB command for hitting enter was performed.
|
||||
PRESS_ENTER = 7
|
||||
|
||||
########### Episode status actions ###########
|
||||
|
||||
# An action used to indicate the desired task has been completed and resets
|
||||
# the environment. This action should also be used in the case that the task
|
||||
# has already been completed and there is nothing to do.
|
||||
# e.g. The task is to turn on the Wi-Fi when it is already on
|
||||
STATUS_TASK_COMPLETE = 10
|
||||
|
||||
# An action used to indicate that desired task is impossible to complete and
|
||||
# resets the environment. This can be a result of many different things
|
||||
# including UI changes, Android version differences, etc.
|
||||
STATUS_TASK_IMPOSSIBLE = 11
|
||||
@@ -0,0 +1,262 @@
|
||||
from typing import List, Optional, Union, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
|
||||
|
||||
class BoxAnnotator:
|
||||
"""
|
||||
A class for drawing bounding boxes on an image using detections provided.
|
||||
|
||||
Attributes:
|
||||
color (Union[Color, ColorPalette]): The color to draw the bounding box,
|
||||
can be a single color or a color palette
|
||||
thickness (int): The thickness of the bounding box lines, default is 2
|
||||
text_color (Color): The color of the text on the bounding box, default is white
|
||||
text_scale (float): The scale of the text on the bounding box, default is 0.5
|
||||
text_thickness (int): The thickness of the text on the bounding box,
|
||||
default is 1
|
||||
text_padding (int): The padding around the text on the bounding box,
|
||||
default is 5
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.DEFAULT,
|
||||
thickness: int = 3, # 1 for seeclick 2 for mind2web and 3 for demo
|
||||
text_color: Color = Color.BLACK,
|
||||
text_scale: float = 0.5, # 0.8 for mobile/web, 0.3 for desktop # 0.4 for mind2web
|
||||
text_thickness: int = 2, #1, # 2 for demo
|
||||
text_padding: int = 10,
|
||||
avoid_overlap: bool = True,
|
||||
):
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.text_color: Color = text_color
|
||||
self.text_scale: float = text_scale
|
||||
self.text_thickness: int = text_thickness
|
||||
self.text_padding: int = text_padding
|
||||
self.avoid_overlap: bool = avoid_overlap
|
||||
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
labels: Optional[List[str]] = None,
|
||||
skip_label: bool = False,
|
||||
image_size: Optional[Tuple[int, int]] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Draws bounding boxes on the frame using the detections provided.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image on which the bounding boxes will be drawn
|
||||
detections (Detections): The detections for which the
|
||||
bounding boxes will be drawn
|
||||
labels (Optional[List[str]]): An optional list of labels
|
||||
corresponding to each detection. If `labels` are not provided,
|
||||
corresponding `class_id` will be used as label.
|
||||
skip_label (bool): Is set to `True`, skips bounding box label annotation.
|
||||
Returns:
|
||||
np.ndarray: The image with the bounding boxes drawn on it
|
||||
|
||||
Example:
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
classes = ['person', ...]
|
||||
image = ...
|
||||
detections = sv.Detections(...)
|
||||
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
labels = [
|
||||
f"{classes[class_id]} {confidence:0.2f}"
|
||||
for _, _, confidence, class_id, _ in detections
|
||||
]
|
||||
annotated_frame = box_annotator.annotate(
|
||||
scene=image.copy(),
|
||||
detections=detections,
|
||||
labels=labels
|
||||
)
|
||||
```
|
||||
"""
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
for i in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[i].astype(int)
|
||||
class_id = (
|
||||
detections.class_id[i] if detections.class_id is not None else None
|
||||
)
|
||||
idx = class_id if class_id is not None else i
|
||||
color = (
|
||||
self.color.by_idx(idx)
|
||||
if isinstance(self.color, ColorPalette)
|
||||
else self.color
|
||||
)
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=(x1, y1),
|
||||
pt2=(x2, y2),
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
if skip_label:
|
||||
continue
|
||||
|
||||
text = (
|
||||
f"{class_id}"
|
||||
if (labels is None or len(detections) != len(labels))
|
||||
else labels[i]
|
||||
)
|
||||
|
||||
text_width, text_height = cv2.getTextSize(
|
||||
text=text,
|
||||
fontFace=font,
|
||||
fontScale=self.text_scale,
|
||||
thickness=self.text_thickness,
|
||||
)[0]
|
||||
|
||||
if not self.avoid_overlap:
|
||||
text_x = x1 + self.text_padding
|
||||
text_y = y1 - self.text_padding
|
||||
|
||||
text_background_x1 = x1
|
||||
text_background_y1 = y1 - 2 * self.text_padding - text_height
|
||||
|
||||
text_background_x2 = x1 + 2 * self.text_padding + text_width
|
||||
text_background_y2 = y1
|
||||
# text_x = x1 - self.text_padding - text_width
|
||||
# text_y = y1 + self.text_padding + text_height
|
||||
# text_background_x1 = x1 - 2 * self.text_padding - text_width
|
||||
# text_background_y1 = y1
|
||||
# text_background_x2 = x1
|
||||
# text_background_y2 = y1 + 2 * self.text_padding + text_height
|
||||
else:
|
||||
text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2 = get_optimal_label_pos(self.text_padding, text_width, text_height, x1, y1, x2, y2, detections, image_size)
|
||||
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=(text_background_x1, text_background_y1),
|
||||
pt2=(text_background_x2, text_background_y2),
|
||||
color=color.as_bgr(),
|
||||
thickness=cv2.FILLED,
|
||||
)
|
||||
# import pdb; pdb.set_trace()
|
||||
box_color = color.as_rgb()
|
||||
luminance = 0.299 * box_color[0] + 0.587 * box_color[1] + 0.114 * box_color[2]
|
||||
text_color = (0,0,0) if luminance > 160 else (255,255,255)
|
||||
cv2.putText(
|
||||
img=scene,
|
||||
text=text,
|
||||
org=(text_x, text_y),
|
||||
fontFace=font,
|
||||
fontScale=self.text_scale,
|
||||
# color=self.text_color.as_rgb(),
|
||||
color=text_color,
|
||||
thickness=self.text_thickness,
|
||||
lineType=cv2.LINE_AA,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
def box_area(box):
|
||||
return (box[2] - box[0]) * (box[3] - box[1])
|
||||
|
||||
def intersection_area(box1, box2):
|
||||
x1 = max(box1[0], box2[0])
|
||||
y1 = max(box1[1], box2[1])
|
||||
x2 = min(box1[2], box2[2])
|
||||
y2 = min(box1[3], box2[3])
|
||||
return max(0, x2 - x1) * max(0, y2 - y1)
|
||||
|
||||
def IoU(box1, box2, return_max=True):
|
||||
intersection = intersection_area(box1, box2)
|
||||
union = box_area(box1) + box_area(box2) - intersection
|
||||
if box_area(box1) > 0 and box_area(box2) > 0:
|
||||
ratio1 = intersection / box_area(box1)
|
||||
ratio2 = intersection / box_area(box2)
|
||||
else:
|
||||
ratio1, ratio2 = 0, 0
|
||||
if return_max:
|
||||
return max(intersection / union, ratio1, ratio2)
|
||||
else:
|
||||
return intersection / union
|
||||
|
||||
|
||||
def get_optimal_label_pos(text_padding, text_width, text_height, x1, y1, x2, y2, detections, image_size):
|
||||
""" check overlap of text and background detection box, and get_optimal_label_pos,
|
||||
pos: str, position of the text, must be one of 'top left', 'top right', 'outer left', 'outer right' TODO: if all are overlapping, return the last one, i.e. outer right
|
||||
Threshold: default to 0.3
|
||||
"""
|
||||
|
||||
def get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size):
|
||||
is_overlap = False
|
||||
for i in range(len(detections)):
|
||||
detection = detections.xyxy[i].astype(int)
|
||||
if IoU([text_background_x1, text_background_y1, text_background_x2, text_background_y2], detection) > 0.3:
|
||||
is_overlap = True
|
||||
break
|
||||
# check if the text is out of the image
|
||||
if text_background_x1 < 0 or text_background_x2 > image_size[0] or text_background_y1 < 0 or text_background_y2 > image_size[1]:
|
||||
is_overlap = True
|
||||
return is_overlap
|
||||
|
||||
# if pos == 'top left':
|
||||
text_x = x1 + text_padding
|
||||
text_y = y1 - text_padding
|
||||
|
||||
text_background_x1 = x1
|
||||
text_background_y1 = y1 - 2 * text_padding - text_height
|
||||
|
||||
text_background_x2 = x1 + 2 * text_padding + text_width
|
||||
text_background_y2 = y1
|
||||
is_overlap = get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size)
|
||||
if not is_overlap:
|
||||
return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2
|
||||
|
||||
# elif pos == 'outer left':
|
||||
text_x = x1 - text_padding - text_width
|
||||
text_y = y1 + text_padding + text_height
|
||||
|
||||
text_background_x1 = x1 - 2 * text_padding - text_width
|
||||
text_background_y1 = y1
|
||||
|
||||
text_background_x2 = x1
|
||||
text_background_y2 = y1 + 2 * text_padding + text_height
|
||||
is_overlap = get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size)
|
||||
if not is_overlap:
|
||||
return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2
|
||||
|
||||
|
||||
# elif pos == 'outer right':
|
||||
text_x = x2 + text_padding
|
||||
text_y = y1 + text_padding + text_height
|
||||
|
||||
text_background_x1 = x2
|
||||
text_background_y1 = y1
|
||||
|
||||
text_background_x2 = x2 + 2 * text_padding + text_width
|
||||
text_background_y2 = y1 + 2 * text_padding + text_height
|
||||
|
||||
is_overlap = get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size)
|
||||
if not is_overlap:
|
||||
return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2
|
||||
|
||||
# elif pos == 'top right':
|
||||
text_x = x2 - text_padding - text_width
|
||||
text_y = y1 - text_padding
|
||||
|
||||
text_background_x1 = x2 - 2 * text_padding - text_width
|
||||
text_background_y1 = y1 - 2 * text_padding - text_height
|
||||
|
||||
text_background_x2 = x2
|
||||
text_background_y2 = y1
|
||||
|
||||
is_overlap = get_is_overlap(detections, text_background_x1, text_background_y1, text_background_x2, text_background_y2, image_size)
|
||||
if not is_overlap:
|
||||
return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2
|
||||
|
||||
return text_x, text_y, text_background_x1, text_background_y1, text_background_x2, text_background_y2
|
||||
@@ -0,0 +1,423 @@
|
||||
# from ultralytics import YOLO
|
||||
import os
|
||||
import io
|
||||
import base64
|
||||
import time
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
import json
|
||||
import requests
|
||||
# utility function
|
||||
import os
|
||||
#from openai import AzureOpenAI
|
||||
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
# %matplotlib inline
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
import time
|
||||
import base64
|
||||
|
||||
import os
|
||||
import ast
|
||||
import torch
|
||||
from typing import Tuple, List
|
||||
from torchvision.ops import box_convert
|
||||
import re
|
||||
from torchvision.transforms import ToPILImage
|
||||
import supervision as sv
|
||||
import torchvision.transforms as T
|
||||
|
||||
cur_path_ = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
def get_caption_model_processor(model_name, model_name_or_path="Salesforce/blip2-opt-2.7b", device=None):
|
||||
if not device:
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
if model_name == "blip2":
|
||||
from transformers import Blip2Processor, Blip2ForConditionalGeneration
|
||||
processor = Blip2Processor.from_pretrained(os.path.join(cur_path_,"blip2-opt-2.7b"))
|
||||
if device == 'cpu':
|
||||
model = Blip2ForConditionalGeneration.from_pretrained(
|
||||
model_name_or_path, device_map=None, torch_dtype=torch.float32
|
||||
)
|
||||
else:
|
||||
model = Blip2ForConditionalGeneration.from_pretrained(
|
||||
model_name_or_path, device_map=None, torch_dtype=torch.float16
|
||||
).to(device)
|
||||
elif model_name == "florence2":
|
||||
from transformers import AutoProcessor, AutoModelForCausalLM
|
||||
processor = AutoProcessor.from_pretrained(os.path.join(cur_path_,"Florence-2-base"), trust_remote_code=True)
|
||||
if device == 'cpu':
|
||||
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch.float32, trust_remote_code=True)
|
||||
else:
|
||||
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch.float16, trust_remote_code=True).to(device)
|
||||
elif model_name=="default":
|
||||
from transformers import AutoProcessor, AutoModelForCausalLM
|
||||
from huggingface_hub import hf_hub_download
|
||||
processor = AutoProcessor.from_pretrained(os.path.join(cur_path_, "Florence-2-base"), trust_remote_code=True)
|
||||
hf_hub_download("microsoft/OmniParser",filename="generation_config.json",subfolder="icon_caption_florence")
|
||||
hf_hub_download("microsoft/OmniParser", filename="config.json",subfolder="icon_caption_florence")
|
||||
model_name_or_path = hf_hub_download("microsoft/OmniParser", filename="model.safetensors",subfolder="icon_caption_florence")
|
||||
cache_dir=os.path.dirname(model_name_or_path)
|
||||
if device == 'cpu':
|
||||
model = AutoModelForCausalLM.from_pretrained(cache_dir,torch_dtype=torch.float32,
|
||||
trust_remote_code=True)
|
||||
else:
|
||||
model = AutoModelForCausalLM.from_pretrained(cache_dir, torch_dtype=torch.float16,
|
||||
trust_remote_code=True).to(device)
|
||||
|
||||
return {'model': model.to(device), 'processor': processor}
|
||||
|
||||
|
||||
def get_yolo_model(model_path):
|
||||
from ultralytics import YOLO
|
||||
# Load the model.
|
||||
model = YOLO(model_path)
|
||||
return model
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def get_parsed_content_icon(filtered_boxes, ocr_bbox, image_source, caption_model_processor, prompt=None):
|
||||
to_pil = ToPILImage()
|
||||
if ocr_bbox:
|
||||
non_ocr_boxes = filtered_boxes[len(ocr_bbox):]
|
||||
else:
|
||||
non_ocr_boxes = filtered_boxes
|
||||
croped_pil_image = []
|
||||
for i, coord in enumerate(non_ocr_boxes):
|
||||
xmin, xmax = int(coord[0]*image_source.shape[1]), int(coord[2]*image_source.shape[1])
|
||||
ymin, ymax = int(coord[1]*image_source.shape[0]), int(coord[3]*image_source.shape[0])
|
||||
cropped_image = image_source[ymin:ymax, xmin:xmax, :]
|
||||
croped_pil_image.append(to_pil(cropped_image))
|
||||
|
||||
model, processor = caption_model_processor['model'], caption_model_processor['processor']
|
||||
if not prompt:
|
||||
if 'florence' in model.config.name_or_path:
|
||||
prompt = "<CAPTION>"
|
||||
else:
|
||||
prompt = "The image shows"
|
||||
|
||||
batch_size = 10 # Number of samples per batch
|
||||
generated_texts = []
|
||||
device = model.device
|
||||
|
||||
for i in range(0, len(croped_pil_image), batch_size):
|
||||
batch = croped_pil_image[i:i+batch_size]
|
||||
if model.device.type == 'cuda':
|
||||
inputs = processor(images=batch, text=[prompt]*len(batch), return_tensors="pt").to(device=device, dtype=torch.float16)
|
||||
else:
|
||||
inputs = processor(images=batch, text=[prompt]*len(batch), return_tensors="pt").to(device=device)
|
||||
if 'florence' in model.config.name_or_path:
|
||||
generated_ids = model.generate(input_ids=inputs["input_ids"],pixel_values=inputs["pixel_values"],max_new_tokens=1024,num_beams=3, do_sample=False)
|
||||
else:
|
||||
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,
|
||||
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
|
||||
generated_text = [gen.strip() for gen in generated_text]
|
||||
generated_texts.extend(generated_text)
|
||||
|
||||
return generated_texts
|
||||
|
||||
|
||||
|
||||
def get_parsed_content_icon_phi3v(filtered_boxes, ocr_bbox, image_source, caption_model_processor):
|
||||
to_pil = ToPILImage()
|
||||
if ocr_bbox:
|
||||
non_ocr_boxes = filtered_boxes[len(ocr_bbox):]
|
||||
else:
|
||||
non_ocr_boxes = filtered_boxes
|
||||
croped_pil_image = []
|
||||
for i, coord in enumerate(non_ocr_boxes):
|
||||
xmin, xmax = int(coord[0]*image_source.shape[1]), int(coord[2]*image_source.shape[1])
|
||||
ymin, ymax = int(coord[1]*image_source.shape[0]), int(coord[3]*image_source.shape[0])
|
||||
cropped_image = image_source[ymin:ymax, xmin:xmax, :]
|
||||
croped_pil_image.append(to_pil(cropped_image))
|
||||
|
||||
model, processor = caption_model_processor['model'], caption_model_processor['processor']
|
||||
device = model.device
|
||||
messages = [{"role": "user", "content": "<|image_1|>\ndescribe the icon in one sentence"}]
|
||||
prompt = processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
||||
|
||||
batch_size = 5 # Number of samples per batch
|
||||
generated_texts = []
|
||||
|
||||
for i in range(0, len(croped_pil_image), batch_size):
|
||||
images = croped_pil_image[i:i+batch_size]
|
||||
image_inputs = [processor.image_processor(x, return_tensors="pt") for x in images]
|
||||
inputs ={'input_ids': [], 'attention_mask': [], 'pixel_values': [], 'image_sizes': []}
|
||||
texts = [prompt] * len(images)
|
||||
for i, txt in enumerate(texts):
|
||||
input = processor._convert_images_texts_to_inputs(image_inputs[i], txt, return_tensors="pt")
|
||||
inputs['input_ids'].append(input['input_ids'])
|
||||
inputs['attention_mask'].append(input['attention_mask'])
|
||||
inputs['pixel_values'].append(input['pixel_values'])
|
||||
inputs['image_sizes'].append(input['image_sizes'])
|
||||
max_len = max([x.shape[1] for x in inputs['input_ids']])
|
||||
for i, v in enumerate(inputs['input_ids']):
|
||||
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)
|
||||
inputs['attention_mask'][i] = torch.cat([torch.zeros(1, max_len - v.shape[1], dtype=torch.long), inputs['attention_mask'][i]], dim=1)
|
||||
inputs_cat = {k: torch.concatenate(v).to(device) for k, v in inputs.items()}
|
||||
|
||||
generation_args = {
|
||||
"max_new_tokens": 25,
|
||||
"temperature": 0.01,
|
||||
"do_sample": False,
|
||||
}
|
||||
generate_ids = model.generate(**inputs_cat, eos_token_id=processor.tokenizer.eos_token_id, **generation_args)
|
||||
# # remove input tokens
|
||||
generate_ids = generate_ids[:, inputs_cat['input_ids'].shape[1]:]
|
||||
response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
|
||||
response = [res.strip('\n').strip() for res in response]
|
||||
generated_texts.extend(response)
|
||||
|
||||
return generated_texts
|
||||
|
||||
def remove_overlap(boxes, iou_threshold, ocr_bbox=None):
|
||||
assert ocr_bbox is None or isinstance(ocr_bbox, List)
|
||||
|
||||
def box_area(box):
|
||||
return (box[2] - box[0]) * (box[3] - box[1])
|
||||
|
||||
def intersection_area(box1, box2):
|
||||
x1 = max(box1[0], box2[0])
|
||||
y1 = max(box1[1], box2[1])
|
||||
x2 = min(box1[2], box2[2])
|
||||
y2 = min(box1[3], box2[3])
|
||||
return max(0, x2 - x1) * max(0, y2 - y1)
|
||||
|
||||
def IoU(box1, box2):
|
||||
intersection = intersection_area(box1, box2)
|
||||
union = box_area(box1) + box_area(box2) - intersection + 1e-6
|
||||
if box_area(box1) > 0 and box_area(box2) > 0:
|
||||
ratio1 = intersection / box_area(box1)
|
||||
ratio2 = intersection / box_area(box2)
|
||||
else:
|
||||
ratio1, ratio2 = 0, 0
|
||||
return max(intersection / union, ratio1, ratio2)
|
||||
|
||||
boxes = boxes.tolist()
|
||||
filtered_boxes = []
|
||||
if ocr_bbox:
|
||||
filtered_boxes.extend(ocr_bbox)
|
||||
# print('ocr_bbox!!!', ocr_bbox)
|
||||
for i, box1 in enumerate(boxes):
|
||||
# if not any(IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2) for j, box2 in enumerate(boxes) if i != j):
|
||||
is_valid_box = True
|
||||
for j, box2 in enumerate(boxes):
|
||||
if i != j and IoU(box1, box2) > iou_threshold and box_area(box1) > box_area(box2):
|
||||
is_valid_box = False
|
||||
break
|
||||
if is_valid_box:
|
||||
# add the following 2 lines to include ocr bbox
|
||||
if ocr_bbox:
|
||||
if not any(IoU(box1, box3) > iou_threshold for k, box3 in enumerate(ocr_bbox)):
|
||||
filtered_boxes.append(box1)
|
||||
else:
|
||||
filtered_boxes.append(box1)
|
||||
return torch.tensor(filtered_boxes)
|
||||
|
||||
def load_image(image_path: str) -> Tuple[np.array, torch.Tensor]:
|
||||
transform = T.Compose(
|
||||
[
|
||||
T.RandomResize([800], max_size=1333),
|
||||
T.ToTensor(),
|
||||
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
||||
]
|
||||
)
|
||||
image_source = Image.open(image_path).convert("RGB")
|
||||
image = np.asarray(image_source)
|
||||
image_transformed, _ = transform(image_source, None)
|
||||
return image, image_transformed
|
||||
|
||||
|
||||
def annotate(image_source: np.ndarray, boxes: torch.Tensor, logits: torch.Tensor, phrases: List[str], text_scale: float,
|
||||
text_padding=5, text_thickness=2, thickness=3) -> np.ndarray:
|
||||
"""
|
||||
This function annotates an image with bounding boxes and labels.
|
||||
|
||||
Parameters:
|
||||
image_source (np.ndarray): The source image to be annotated.
|
||||
boxes (torch.Tensor): A tensor containing bounding box coordinates. in cxcywh format, pixel scale
|
||||
logits (torch.Tensor): A tensor containing confidence scores for each bounding box.
|
||||
phrases (List[str]): A list of labels for each bounding box.
|
||||
text_scale (float): The scale of the text to be displayed. 0.8 for mobile/web, 0.3 for desktop # 0.4 for mind2web
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
"""
|
||||
h, w, _ = image_source.shape
|
||||
boxes = boxes * torch.Tensor([w, h, w, h])
|
||||
xyxy = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy()
|
||||
xywh = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xywh").numpy()
|
||||
detections = sv.Detections(xyxy=xyxy)
|
||||
|
||||
labels = [f"{phrase}" for phrase in range(boxes.shape[0])]
|
||||
|
||||
from .util.box_annotator import BoxAnnotator
|
||||
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
|
||||
annotated_frame = image_source.copy()
|
||||
annotated_frame = box_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels, image_size=(w,h))
|
||||
|
||||
label_coordinates = {f"{phrase}": v for phrase, v in zip(phrases, xywh)}
|
||||
return annotated_frame, label_coordinates
|
||||
|
||||
|
||||
def predict(model, image, caption, box_threshold, text_threshold):
|
||||
""" Use huggingface model to replace the original model
|
||||
"""
|
||||
model, processor = model['model'], model['processor']
|
||||
device = model.device
|
||||
|
||||
inputs = processor(images=image, text=caption, return_tensors="pt").to(device)
|
||||
with torch.no_grad():
|
||||
outputs = model(**inputs)
|
||||
|
||||
results = processor.post_process_grounded_object_detection(
|
||||
outputs,
|
||||
inputs.input_ids,
|
||||
box_threshold=box_threshold, # 0.4,
|
||||
text_threshold=text_threshold, # 0.3,
|
||||
target_sizes=[image.size[::-1]]
|
||||
)[0]
|
||||
boxes, logits, phrases = results["boxes"], results["scores"], results["labels"]
|
||||
return boxes, logits, phrases
|
||||
|
||||
|
||||
def predict_yolo(model, image_path, box_threshold):
|
||||
""" Use huggingface model to replace the original model
|
||||
"""
|
||||
# model = model['model']
|
||||
|
||||
result = model.predict(
|
||||
source=image_path,
|
||||
conf=box_threshold,
|
||||
# iou=0.5, # default 0.7
|
||||
)
|
||||
boxes = result[0].boxes.xyxy#.tolist() # in pixel space
|
||||
conf = result[0].boxes.conf
|
||||
phrases = [str(i) for i in range(len(boxes))]
|
||||
|
||||
return boxes, conf, phrases
|
||||
|
||||
|
||||
def get_som_labeled_img(optional_c,img_path, model=None,batch_num=0, 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):
|
||||
""" ocr_bbox: list of xyxy format bbox
|
||||
"""
|
||||
TEXT_PROMPT = "clickable buttons on the screen"
|
||||
# BOX_TRESHOLD = 0.02 # 0.05/0.02 for web and 0.1 for mobile
|
||||
TEXT_TRESHOLD = 0.01 # 0.9 # 0.01
|
||||
image_source = Image.open(img_path).convert("RGB")
|
||||
w, h = image_source.size
|
||||
# import pdb; pdb.set_trace()
|
||||
if not optional_c: # False: # TODO
|
||||
xyxy, logits, phrases = predict(model=model, image=image_source, caption=TEXT_PROMPT, box_threshold=BOX_TRESHOLD, text_threshold=TEXT_TRESHOLD)
|
||||
else:
|
||||
xyxy, logits, phrases = predict_yolo(model=model, image_path=img_path, box_threshold=BOX_TRESHOLD)
|
||||
xyxy = xyxy / torch.Tensor([w, h, w, h]).to(xyxy.device)
|
||||
image_source = np.asarray(image_source)
|
||||
phrases = [str(i) for i in range(len(phrases))]
|
||||
|
||||
# annotate the image with labels
|
||||
h, w, _ = image_source.shape
|
||||
if ocr_bbox:
|
||||
ocr_bbox = torch.tensor(ocr_bbox) / torch.Tensor([w, h, w, h])
|
||||
ocr_bbox=ocr_bbox.tolist()
|
||||
else:
|
||||
print('no ocr bbox!!!')
|
||||
ocr_bbox = None
|
||||
filtered_boxes = remove_overlap(boxes=xyxy, iou_threshold=iou_threshold, ocr_bbox=ocr_bbox)
|
||||
|
||||
# get parsed icon local semantics
|
||||
if use_local_semantics:
|
||||
caption_model = caption_model_processor['model']
|
||||
if 'phi3_v' in caption_model.config.model_type:
|
||||
parsed_content_icon = get_parsed_content_icon_phi3v(filtered_boxes, ocr_bbox, image_source, caption_model_processor)
|
||||
else:
|
||||
parsed_content_icon = get_parsed_content_icon(filtered_boxes, ocr_bbox, image_source, caption_model_processor, prompt=prompt)
|
||||
ocr_text = [f"IMG{batch_num}'s Text Box ID {i}: {txt}" for i, txt in enumerate(ocr_text)]
|
||||
icon_start = len(ocr_text)
|
||||
parsed_content_icon_ls = []
|
||||
for i, txt in enumerate(parsed_content_icon):
|
||||
parsed_content_icon_ls.append(f"IMG{batch_num}'s Icon Box ID {str(i+icon_start)}: {txt}")
|
||||
parsed_content_merged = ocr_text + parsed_content_icon_ls
|
||||
else:
|
||||
ocr_text = [f"IMG{batch_num}'s Text Box ID {i}: {txt}" for i, txt in enumerate(ocr_text)]
|
||||
parsed_content_merged = ocr_text
|
||||
|
||||
filtered_boxes = box_convert(boxes=filtered_boxes, in_fmt="xyxy", out_fmt="cxcywh")
|
||||
|
||||
phrases = [i for i in range(len(filtered_boxes))]
|
||||
|
||||
# draw boxes
|
||||
if draw_bbox_config:
|
||||
annotated_frame, label_coordinates = annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, **draw_bbox_config)
|
||||
else:
|
||||
annotated_frame, label_coordinates = annotate(image_source=image_source, boxes=filtered_boxes, logits=logits, phrases=phrases, text_scale=text_scale, text_padding=text_padding)
|
||||
|
||||
pil_img = Image.fromarray(annotated_frame)
|
||||
buffered = io.BytesIO()
|
||||
pil_img.save(buffered, format="PNG")
|
||||
encoded_image = base64.b64encode(buffered.getvalue()).decode('ascii')
|
||||
if output_coord_in_ratio:
|
||||
# h, w, _ = image_source.shape
|
||||
label_coordinates = {k: [v[0]/w, v[1]/h, v[2]/w, v[3]/h] for k, v in label_coordinates.items()}
|
||||
assert w == annotated_frame.shape[1] and h == annotated_frame.shape[0]
|
||||
|
||||
return encoded_image, label_coordinates, parsed_content_merged
|
||||
|
||||
|
||||
def get_xywh(input):
|
||||
x, y, w, h = input[0][0], input[0][1], input[2][0] - input[0][0], input[2][1] - input[0][1]
|
||||
x, y, w, h = int(x), int(y), int(w), int(h)
|
||||
return x, y, w, h
|
||||
|
||||
def get_xyxy(input):
|
||||
x, y, xp, yp = input[0][0], input[0][1], input[2][0], input[2][1]
|
||||
x, y, xp, yp = int(x), int(y), int(xp), int(yp)
|
||||
return x, y, xp, yp
|
||||
|
||||
def get_xywh_yolo(input):
|
||||
x, y, w, h = input[0], input[1], input[2] - input[0], input[3] - input[1]
|
||||
x, y, w, h = int(x), int(y), int(w), int(h)
|
||||
return x, y, w, h
|
||||
|
||||
|
||||
|
||||
def check_ocr_box(image_path,ORC, use_paddleocr = False,display_img = True, output_bb_format='xywh', goal_filtering=None, easyocr_args=None,):
|
||||
is_goal_filtered = False
|
||||
if not use_paddleocr:
|
||||
if easyocr_args is None:
|
||||
easyocr_args = {}
|
||||
result = ORC.readtext(image_path, **easyocr_args)
|
||||
# print('goal filtering pred:', result[-5:])
|
||||
coord = [item[0] for item in result]
|
||||
text = [item[1] for item in result]
|
||||
else:
|
||||
result = ORC.ocr(image_path, cls=False)[0]
|
||||
coord = [item[0] for item in result]
|
||||
text = [item[1][0] for item in result]
|
||||
# read the image using cv2
|
||||
if display_img:
|
||||
opencv_img = cv2.imread(image_path)
|
||||
opencv_img = cv2.cvtColor(opencv_img, cv2.COLOR_RGB2BGR)
|
||||
bb = []
|
||||
for item in coord:
|
||||
x, y, a, b = get_xywh(item)
|
||||
# print(x, y, a, b)
|
||||
bb.append((x, y, a, b))
|
||||
cv2.rectangle(opencv_img, (x, y), (x+a, y+b), (0, 255, 0), 2)
|
||||
|
||||
# Display the image
|
||||
plt.imshow(opencv_img)
|
||||
else:
|
||||
if output_bb_format == 'xywh':
|
||||
bb = [get_xywh(item) for item in coord]
|
||||
elif output_bb_format == 'xyxy':
|
||||
bb = [get_xyxy(item) for item in coord]
|
||||
# print('bounding box!!!', bb)
|
||||
return (text, bb), is_goal_filtered
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
import torch
|
||||
from ultralytics.nn.tasks import DetectionModel
|
||||
from safetensors.torch import load_file
|
||||
|
||||
def convert_safetensor_to_pt():
|
||||
tensor_dict = load_file("weights/icon_detect/model.safetensors")
|
||||
|
||||
model = DetectionModel('weights/icon_detect/model.yaml')
|
||||
model.load_state_dict(tensor_dict)
|
||||
torch.save({'model': model}, 'weights/icon_detect/best.pt')
|
||||
|
||||
|
||||
Binary file not shown.
@@ -0,0 +1,127 @@
|
||||
backbone:
|
||||
- - -1
|
||||
- 1
|
||||
- Conv
|
||||
- - 64
|
||||
- 3
|
||||
- 2
|
||||
- - -1
|
||||
- 1
|
||||
- Conv
|
||||
- - 128
|
||||
- 3
|
||||
- 2
|
||||
- - -1
|
||||
- 3
|
||||
- C2f
|
||||
- - 128
|
||||
- true
|
||||
- - -1
|
||||
- 1
|
||||
- Conv
|
||||
- - 256
|
||||
- 3
|
||||
- 2
|
||||
- - -1
|
||||
- 6
|
||||
- C2f
|
||||
- - 256
|
||||
- true
|
||||
- - -1
|
||||
- 1
|
||||
- Conv
|
||||
- - 512
|
||||
- 3
|
||||
- 2
|
||||
- - -1
|
||||
- 6
|
||||
- C2f
|
||||
- - 512
|
||||
- true
|
||||
- - -1
|
||||
- 1
|
||||
- Conv
|
||||
- - 1024
|
||||
- 3
|
||||
- 2
|
||||
- - -1
|
||||
- 3
|
||||
- C2f
|
||||
- - 1024
|
||||
- true
|
||||
- - -1
|
||||
- 1
|
||||
- SPPF
|
||||
- - 1024
|
||||
- 5
|
||||
ch: 3
|
||||
depth_multiple: 0.33
|
||||
head:
|
||||
- - -1
|
||||
- 1
|
||||
- nn.Upsample
|
||||
- - None
|
||||
- 2
|
||||
- nearest
|
||||
- - - -1
|
||||
- 6
|
||||
- 1
|
||||
- Concat
|
||||
- - 1
|
||||
- - -1
|
||||
- 3
|
||||
- C2f
|
||||
- - 512
|
||||
- - -1
|
||||
- 1
|
||||
- nn.Upsample
|
||||
- - None
|
||||
- 2
|
||||
- nearest
|
||||
- - - -1
|
||||
- 4
|
||||
- 1
|
||||
- Concat
|
||||
- - 1
|
||||
- - -1
|
||||
- 3
|
||||
- C2f
|
||||
- - 256
|
||||
- - -1
|
||||
- 1
|
||||
- Conv
|
||||
- - 256
|
||||
- 3
|
||||
- 2
|
||||
- - - -1
|
||||
- 12
|
||||
- 1
|
||||
- Concat
|
||||
- - 1
|
||||
- - -1
|
||||
- 3
|
||||
- C2f
|
||||
- - 512
|
||||
- - -1
|
||||
- 1
|
||||
- Conv
|
||||
- - 512
|
||||
- 3
|
||||
- 2
|
||||
- - - -1
|
||||
- 9
|
||||
- 1
|
||||
- Concat
|
||||
- - 1
|
||||
- - -1
|
||||
- 3
|
||||
- C2f
|
||||
- - 1024
|
||||
- - - 15
|
||||
- 18
|
||||
- 21
|
||||
- 1
|
||||
- Detect
|
||||
- - nc
|
||||
nc: 1
|
||||
width_multiple: 0.25
|
||||
@@ -0,0 +1,189 @@
|
||||
# !/usr/bin/env python
|
||||
# -*- coding: UTF-8 -*-
|
||||
import base64, os
|
||||
import io
|
||||
import random
|
||||
from PIL import Image
|
||||
import torch
|
||||
import logging
|
||||
import easyocr
|
||||
|
||||
import folder_paths
|
||||
from .utils import get_config,load_list_images,convert_safetensor_to_pt,tensor2imglist
|
||||
from .OmniParser.utils import check_ocr_box, get_yolo_model, get_caption_model_processor, get_som_labeled_img
|
||||
|
||||
node_cur_path = os.path.dirname(os.path.abspath(__file__))
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
class OmniParser_Loader:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"repo": ("STRING", {"default": "microsoft/OmniParser", }),
|
||||
"platform": (["pc", "web","mobile"],),
|
||||
},
|
||||
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("OP_MODEL",)
|
||||
RETURN_NAMES = ("model",)
|
||||
FUNCTION = "main"
|
||||
CATEGORY = "OmniParser"
|
||||
|
||||
def main(self, repo,platform,):
|
||||
if repo:
|
||||
if "florence" in repo.lower():
|
||||
model_name="florence2"
|
||||
elif "blip" in repo.lower():
|
||||
model_name="blip2"
|
||||
elif repo== "microsoft/OmniParser":
|
||||
model_name = "default"
|
||||
else:
|
||||
raise "unsupport repo id"
|
||||
else:
|
||||
raise "must fill repo or local path"
|
||||
|
||||
yolo_model_converted=os.path.join(node_cur_path,"OmniParser/weights/icon_detect/best.pt")
|
||||
if not os.path.exists(yolo_model_converted):
|
||||
convert_safetensor_to_pt()
|
||||
logging.info("convert_safetensor_to_pt is done.")
|
||||
|
||||
yolo_model = get_yolo_model(model_path=yolo_model_converted)
|
||||
|
||||
caption_model_processor = get_caption_model_processor(model_name=model_name,
|
||||
model_name_or_path=repo)
|
||||
draw_bbox_config=get_config(platform)
|
||||
|
||||
logging.info("loading checkpoint done.")
|
||||
model={"yolo_model":yolo_model,"caption_model_processor":caption_model_processor,"draw_bbox_config":draw_bbox_config}
|
||||
return (model,)
|
||||
|
||||
class OmniParser_Sampler:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model":("OP_MODEL",),
|
||||
"image": ("IMAGE",), # [B,H,W,C], C=3
|
||||
"text_threshold":("FLOAT", {
|
||||
"default": 0.9,
|
||||
"min": 0.1,
|
||||
"max": 1.0,
|
||||
"step": 0.1,
|
||||
"display": "number",}),
|
||||
"box_threshold": ("FLOAT", {
|
||||
"default": 0.05,
|
||||
"min": 0.01,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"display": "number", }),
|
||||
"iou_threshold": ("FLOAT", {
|
||||
"default": 0.1,
|
||||
"min": 0.01,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"display": "number", }),
|
||||
"orc_model": (["easyocr", "paddleocr"],),
|
||||
"orc_mode": (["en", "'ch_sim','en'"],),
|
||||
"output_bb_format": (["xyxy", "xywh"],),
|
||||
"display_img": ("BOOLEAN", {"default": False},),
|
||||
"paragraph": ("BOOLEAN", {"default": False},),
|
||||
"output_coord_in_ratio": ("BOOLEAN", {"default": True},),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE","STRING")
|
||||
RETURN_NAMES = ("image","string")
|
||||
FUNCTION = "main"
|
||||
CATEGORY = "OmniParser"
|
||||
|
||||
def main(self, model,image,text_threshold,box_threshold,iou_threshold,orc_model,orc_mode,output_bb_format,display_img,paragraph,output_coord_in_ratio,**kwargs):
|
||||
#optional chocie
|
||||
optional_c=True
|
||||
|
||||
#pre model
|
||||
yolo_model=model.get("yolo_model")
|
||||
draw_bbox_config=model.get("draw_bbox_config")
|
||||
caption_model_processor=model.get("caption_model_processor")
|
||||
|
||||
if orc_model=="easyocr":
|
||||
ORC = easyocr.Reader([orc_mode])
|
||||
use_paddleocr = False
|
||||
else:
|
||||
use_paddleocr = True #another choice
|
||||
from paddleocr import PaddleOCR
|
||||
if orc_mode=="'ch_sim','en'":
|
||||
orc_mode="ch"
|
||||
ORC = PaddleOCR(
|
||||
lang=orc_mode, # other lang also available
|
||||
use_angle_cls=False,
|
||||
use_gpu=False, # using cuda will conflict with pytorch in the same process
|
||||
show_log=False,
|
||||
max_batch_size=1024,
|
||||
use_dilation=True, # improves accuracy
|
||||
det_db_score_mode='slow', # improves accuracy
|
||||
rec_batch_num=1024)
|
||||
|
||||
#pre image if list
|
||||
file_prefix = ''.join(random.choice("0123456789abcdefg") for _ in range(6))
|
||||
image_list,B = tensor2imglist(image, np_out=True)
|
||||
path_list = []
|
||||
if B==1:
|
||||
image_input = image_list[0]
|
||||
image_save_path = os.path.join(folder_paths.get_input_directory(), f'Temp_{file_prefix}.png')
|
||||
image_input.save(image_save_path)
|
||||
path_list=[image_save_path]
|
||||
else:
|
||||
for i,img in enumerate(image_list):
|
||||
image_save_path = os.path.join(folder_paths.get_input_directory(), f'Temp_{file_prefix}_{i}.png')
|
||||
img.save(image_save_path)
|
||||
path_list.append(image_save_path)
|
||||
|
||||
ouput_img=[]
|
||||
ouput_text=[]
|
||||
for batch_num,image_path in enumerate(path_list):
|
||||
# pre box
|
||||
ocr_bbox_rslt, is_goal_filtered = check_ocr_box(image_path, ORC, use_paddleocr,
|
||||
display_img=display_img, output_bb_format=output_bb_format,
|
||||
goal_filtering=None,
|
||||
easyocr_args={'paragraph': paragraph,
|
||||
'text_threshold': text_threshold})
|
||||
text, ocr_bbox = ocr_bbox_rslt
|
||||
# infer
|
||||
logging.info("start infer..")
|
||||
dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(optional_c, image_path,
|
||||
yolo_model,batch_num,
|
||||
BOX_TRESHOLD=box_threshold,
|
||||
output_coord_in_ratio=output_coord_in_ratio,
|
||||
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)))
|
||||
ouput_img.append(image)
|
||||
content_list = '\n'.join(parsed_content_list)
|
||||
ouput_text.append(content_list)
|
||||
print('finish processing')
|
||||
image=load_list_images(ouput_img)
|
||||
prompt='\n\n'.join(ouput_text)
|
||||
return (image,prompt)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"OmniParser_Loader": OmniParser_Loader,
|
||||
"OmniParser_Sampler":OmniParser_Sampler,
|
||||
}
|
||||
|
||||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"OmniParser_Loader": "OmniParser_Loader",
|
||||
"OmniParser_Sampler":"OmniParser_Sampler",
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
|
||||
from .OmniParser_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -0,0 +1,15 @@
|
||||
[project]
|
||||
name = "comfyui_omniparser"
|
||||
description = "Try OmniParser in ComfyUI which a simple screen parsing tool towards pure vision based GUI agent."
|
||||
version = "1.0.0"
|
||||
license = "LICENSE"
|
||||
dependencies = ["#torch", "easyocr", "#torchvision", "#supervision==0.18.0", "supervision", "#openai==1.3.5", "#transformers", "#ultralytics==8.1.24", "#azure-identity", "#numpy", "opencv-python", "opencv-python-headless", "#gradio", "dill", "#accelerate", "#timm", "#einops==0.8.0", "#paddlepaddle", "#paddleocr"]
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/smthemex/ComfyUI_OmniParser"
|
||||
# Used by Comfy Registry https://comfyregistry.org
|
||||
|
||||
[tool.comfy]
|
||||
PublisherId = "smthemex"
|
||||
DisplayName = "ComfyUI_OmniParser"
|
||||
Icon = ""
|
||||
@@ -0,0 +1,19 @@
|
||||
#torch
|
||||
easyocr
|
||||
#torchvision
|
||||
#supervision==0.18.0
|
||||
supervision
|
||||
#openai==1.3.5
|
||||
#transformers
|
||||
#ultralytics==8.1.24
|
||||
#azure-identity
|
||||
#numpy
|
||||
opencv-python
|
||||
opencv-python-headless
|
||||
#gradio
|
||||
dill
|
||||
#accelerate
|
||||
#timm
|
||||
#einops==0.8.0
|
||||
#paddlepaddle
|
||||
#paddleocr
|
||||
@@ -0,0 +1,196 @@
|
||||
# !/usr/bin/env python
|
||||
# -*- coding: UTF-8 -*-
|
||||
from ultralytics.nn.tasks import DetectionModel
|
||||
from safetensors.torch import load_file
|
||||
from PIL import Image, ImageOps
|
||||
import numpy as np
|
||||
import cv2
|
||||
import os
|
||||
import torch
|
||||
from comfy.utils import common_upscale,ProgressBar
|
||||
import folder_paths
|
||||
|
||||
cur_path = os.path.dirname(os.path.abspath(__file__))
|
||||
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
||||
|
||||
POINT_SIZE = 4 # Size of the query point in the preview video
|
||||
|
||||
|
||||
def convert_safetensor_to_pt():
|
||||
tensor_dict = load_file(os.path.join(cur_path,"OmniParser/weights/icon_detect/model.safetensors"))
|
||||
|
||||
model = DetectionModel(os.path.join(cur_path,'OmniParser/weights/icon_detect/model.yaml'))
|
||||
model.load_state_dict(tensor_dict)
|
||||
torch.save({'model': model}, os.path.join(cur_path,'OmniParser/weights/icon_detect/best.pt'))
|
||||
|
||||
|
||||
def get_config(platform="pc"):
|
||||
#platform = 'pc'
|
||||
if platform == 'pc':
|
||||
draw_bbox_config = {
|
||||
'text_scale': 0.8,
|
||||
'text_thickness': 2,
|
||||
'text_padding': 2,
|
||||
'thickness': 2,
|
||||
}
|
||||
elif platform == 'web':
|
||||
draw_bbox_config = {
|
||||
'text_scale': 0.8,
|
||||
'text_thickness': 2,
|
||||
'text_padding': 3,
|
||||
'thickness': 3,
|
||||
}
|
||||
elif platform == 'mobile':
|
||||
draw_bbox_config = {
|
||||
'text_scale': 0.8,
|
||||
'text_thickness': 2,
|
||||
'text_padding': 3,
|
||||
'thickness': 3,
|
||||
}
|
||||
else:
|
||||
raise "error"
|
||||
return draw_bbox_config
|
||||
|
||||
def spilit_tensor2list(img_tensor):#[B,H,W,C], C=3,B>=1
|
||||
video_list = []
|
||||
if isinstance(img_tensor, list):
|
||||
if isinstance(img_tensor[0], torch.Tensor):
|
||||
video_list = img_tensor
|
||||
elif isinstance(img_tensor, torch.Tensor):
|
||||
b, _, _, _ = img_tensor.size()
|
||||
if b == 1:
|
||||
img = [b]
|
||||
while img is not []:
|
||||
video_list += img
|
||||
else:
|
||||
video_list = torch.chunk(img_tensor, chunks=b)
|
||||
return video_list
|
||||
|
||||
def tensor2imglist(image,np_out=True):# pil first
|
||||
B, _, _, _ = image.size()
|
||||
if B == 1:
|
||||
if np_out:
|
||||
list_out = [tensor2pil(image)]
|
||||
else:
|
||||
list_out = [tensor2cv(image.squeeze())]
|
||||
else:
|
||||
image_list = torch.chunk(image, chunks=B)
|
||||
if np_out:
|
||||
list_out = [tensor2pil(i) for i in image_list]
|
||||
else:
|
||||
list_out = [tensor2cv(i.squeeze()) for i in image_list]
|
||||
return list_out,B
|
||||
|
||||
|
||||
def tensor2pil(tensor):
|
||||
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
||||
image = Image.fromarray(image_np, mode='RGB')
|
||||
return image
|
||||
|
||||
def pil2narry(img):
|
||||
narry = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
|
||||
return narry
|
||||
|
||||
def tensor_upscale2pil(img_tensor, width, height):
|
||||
samples = img_tensor.movedim(-1, 1)
|
||||
img = common_upscale(samples, width, height, "nearest-exact", "center")
|
||||
samples = img.movedim(1, -1)
|
||||
img_pil = tensor2pil(samples)
|
||||
return img_pil
|
||||
|
||||
def tensor_upscale(img_tensor, width, height): #torch tensor
|
||||
samples = img_tensor.movedim(-1, 1)
|
||||
img = common_upscale(samples, width, height, "nearest-exact", "center")
|
||||
samples = img.movedim(1, -1)
|
||||
return samples
|
||||
|
||||
def convert_from_cv2_to_image(img: np.ndarray) -> Image:
|
||||
# return Image.fromarray(img)
|
||||
return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
|
||||
|
||||
|
||||
def convert_from_image_to_cv2(img: Image) -> np.ndarray:
|
||||
# return np.asarray(img)
|
||||
return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
|
||||
|
||||
def tensor2cv(tensor_image):
|
||||
if tensor_image.is_cuda:
|
||||
tensor_image = tensor_image.detach().cpu()
|
||||
tensor_image=tensor_image.numpy()
|
||||
#反归一化
|
||||
maxValue=tensor_image.max()
|
||||
tensor_image=tensor_image*255/maxValue
|
||||
img_cv2=np.uint8(tensor_image)#32 to uint8
|
||||
img_cv2=cv2.cvtColor(img_cv2,cv2.COLOR_RGB2BGR)
|
||||
return img_cv2
|
||||
|
||||
def cvargb2tensor(img):
|
||||
assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img))
|
||||
img = torch.from_numpy(img.transpose((2, 0, 1)))
|
||||
return img.float().div(255).unsqueeze(0) # 255也可以改为256
|
||||
|
||||
def cv2tensor(img):
|
||||
assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img))
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
img = torch.from_numpy(img.transpose((2, 0, 1)))
|
||||
return img.float().div(255).unsqueeze(0) # 255也可以改为256
|
||||
|
||||
def images_generator(img_list: list,):
|
||||
#get img size
|
||||
sizes = {}
|
||||
for image_ in img_list:
|
||||
if isinstance(image_,Image.Image):
|
||||
count = sizes.get(image_.size, 0)
|
||||
sizes[image_.size] = count + 1
|
||||
elif isinstance(image_,np.ndarray):
|
||||
count = sizes.get(image_.shape[:2][::-1], 0)
|
||||
sizes[image_.shape[:2][::-1]] = count + 1
|
||||
else:
|
||||
raise "unsupport image list,must be pil or cv2!!!"
|
||||
size = max(sizes.items(), key=lambda x: x[1])[0]
|
||||
yield size[0], size[1]
|
||||
|
||||
# any to tensor
|
||||
def load_image(img_in):
|
||||
if isinstance(img_in, Image.Image):
|
||||
img_in=img_in.convert("RGB")
|
||||
i = np.array(img_in, dtype=np.float32)
|
||||
i = torch.from_numpy(i).div_(255)
|
||||
if i.shape[0] != size[1] or i.shape[1] != size[0]:
|
||||
i = torch.from_numpy(i).movedim(-1, 0).unsqueeze(0)
|
||||
i = common_upscale(i, size[0], size[1], "lanczos", "center")
|
||||
i = i.squeeze(0).movedim(0, -1).numpy()
|
||||
return i
|
||||
elif isinstance(img_in,np.ndarray):
|
||||
i=cv2.cvtColor(img_in,cv2.COLOR_BGR2RGB).astype(np.float32)
|
||||
i = torch.from_numpy(i).div_(255)
|
||||
print(i.shape)
|
||||
return i
|
||||
else:
|
||||
raise "unsupport image list,must be pil,cv2 or tensor!!!"
|
||||
|
||||
total_images = len(img_list)
|
||||
processed_images = 0
|
||||
pbar = ProgressBar(total_images)
|
||||
images = map(load_image, img_list)
|
||||
try:
|
||||
prev_image = next(images)
|
||||
while True:
|
||||
next_image = next(images)
|
||||
yield prev_image
|
||||
processed_images += 1
|
||||
pbar.update_absolute(processed_images, total_images)
|
||||
prev_image = next_image
|
||||
except StopIteration:
|
||||
pass
|
||||
if prev_image is not None:
|
||||
yield prev_image
|
||||
|
||||
def load_list_images(img_list: list,):
|
||||
gen = images_generator(img_list)
|
||||
(width, height) = next(gen)
|
||||
images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height, width, 3)))))
|
||||
if len(images) == 0:
|
||||
raise FileNotFoundError(f"No images could be loaded .")
|
||||
return images
|
||||
|
||||
Reference in New Issue
Block a user