add florence model; update demo.ipynb

1openwindow · Oct 25, 2024 · cafc5ef · cafc5ef
1 parent b510e61
commit cafc5ef
Show file tree

Hide file tree

Showing 15 changed files with 186 additions and 242 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1 +1,2 @@
-weights/
+weights/icon_caption_blip2
+weights/icon_caption_florence
diff --git a/__pycache__/utils.cpython-312.pyc b/__pycache__/utils.cpython-312.pyc
diff --git a/demo.ipynb b/demo.ipynb
diff --git a/imgs/google_page.png b/imgs/google_page.png
diff --git a/imgs/mobile_4.png b/imgs/mobile_4.png
diff --git a/imgs/pc_1.png b/imgs/pc_1.png
diff --git a/imgs/saved_image_demo.png b/imgs/saved_image_demo.png
diff --git a/imgs/settings.png b/imgs/settings.png
diff --git a/imgs/windows_home.png b/imgs/windows_home.png
diff --git a/imgs/windows_multitab.png b/imgs/windows_multitab.png
diff --git a/requirement.txt b/requirement.txt
@@ -12,4 +12,4 @@ opencv-python-headless
 gradio
 dill
 accelerate
-
+timm
diff --git a/util/__pycache__/__init__.cpython-312.pyc b/util/__pycache__/__init__.cpython-312.pyc
diff --git a/util/__pycache__/box_annotator.cpython-312.pyc b/util/__pycache__/box_annotator.cpython-312.pyc
diff --git a/utils.py b/utils.py
@@ -33,19 +33,27 @@
 import torchvision.transforms as T
 
 
-def get_caption_model_processor(model_name_or_path="Salesforce/blip2-opt-2.7b", device=None):
+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"
-    from transformers import Blip2Processor, Blip2ForConditionalGeneration
-    processor = Blip2Processor.from_pretrained("Salesforce/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
-    )
+    if model_name == "blip2":
+        from transformers import Blip2Processor, Blip2ForConditionalGeneration
+        processor = Blip2Processor.from_pretrained("Salesforce/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("microsoft/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)
     return {'model': model.to(device), 'processor': processor}
 
 

diff --git a/weights/icon_detect/best.pt b/weights/icon_detect/best.pt
-Original file line number
+Diff line change
@@ Expand Up / @@ -12,4 +12,4 @@ opencv-python-headless @@
     gradio
     dill
     accelerate
+    timm