Reset

.gitignore

@@ -1,4 +1,5 @@
 cache/
+.idea
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

.pre-commit-config.yaml

@@ -1,4 +1,3 @@
-exclude: patch
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
   rev: v4.2.0

app.py

@@ -5,14 +5,8 @@ from __future__ import annotations
 import os
 import shutil
 import subprocess
-
 import gradio as gr
 
-if os.getenv('SYSTEM') == 'spaces':
-    with open('patch') as f:
-        subprocess.run('patch -p1'.split(), cwd='prismer', stdin=f)
-    shutil.copytree('prismer/helpers/images', 'prismer/images', dirs_exist_ok=True)
-
 from app_caption import create_demo as create_demo_caption
 from prismer_model import build_deformable_conv, download_models
 
@@ -32,7 +26,7 @@ if (SPACE_ID := os.getenv('SPACE_ID')) is not None:
     description += f'For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings. <a href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a>'
 
 
-with gr.Blocks(css='style.css') as demo:
+with gr.Blocks() as demo:
     gr.Markdown(description)
     with gr.Tabs():
         with gr.TabItem('Zero-shot Image Captioning'):

patch

@@ -1,82 +0,0 @@
-diff --git a/dataset/caption_dataset.py b/dataset/caption_dataset.py
-index 266fdda..0cc5d3f 100644
---- a/dataset/caption_dataset.py
-+++ b/dataset/caption_dataset.py
-@@ -50,7 +50,7 @@ class Caption(Dataset):
-         elif self.dataset == 'demo':
-             img_path_split = self.data_list[index]['image'].split('/')
-             img_name = img_path_split[-2] + '/' + img_path_split[-1]
--            image, labels, labels_info = get_expert_labels('', self.label_path, img_name, 'helpers', self.experts)
-+            image, labels, labels_info = get_expert_labels('prismer', self.label_path, img_name, 'helpers', self.experts)
- 
-         experts = self.transform(image, labels)
-         experts = post_label_process(experts, labels_info)
-diff --git a/dataset/utils.py b/dataset/utils.py
-index b368aac..418358c 100644
---- a/dataset/utils.py
-+++ b/dataset/utils.py
-@@ -5,6 +5,7 @@
- # https://github.com/NVlabs/prismer/blob/main/LICENSE
- 
- import os
-+import pathlib
- import re
- import json
- import torch
-@@ -14,10 +15,12 @@ import torchvision.transforms as transforms
- import torchvision.transforms.functional as transforms_f
- from dataset.randaugment import RandAugment
- 
--COCO_FEATURES = torch.load('dataset/coco_features.pt')['features']
--ADE_FEATURES = torch.load('dataset/ade_features.pt')['features']
--DETECTION_FEATURES = torch.load('dataset/detection_features.pt')['features']
--BACKGROUND_FEATURES = torch.load('dataset/background_features.pt')
-+cur_dir = pathlib.Path(__file__).parent
-+
-+COCO_FEATURES = torch.load(cur_dir / 'coco_features.pt')['features']
-+ADE_FEATURES = torch.load(cur_dir / 'ade_features.pt')['features']
-+DETECTION_FEATURES = torch.load(cur_dir / 'detection_features.pt')['features']
-+BACKGROUND_FEATURES = torch.load(cur_dir / 'background_features.pt')
- 
- 
- class Transform:
-diff --git a/model/prismer.py b/model/prismer.py
-index 080253a..02362a4 100644
---- a/model/prismer.py
-+++ b/model/prismer.py
-@@ -5,6 +5,7 @@
- # https://github.com/NVlabs/prismer/blob/main/LICENSE
- 
- import json
-+import pathlib
- import torch.nn as nn
- 
- from model.modules.vit import load_encoder
-@@ -12,6 +13,9 @@ from model.modules.roberta import load_decoder
- from transformers import RobertaTokenizer, RobertaConfig
- 
- 
-+cur_dir = pathlib.Path(__file__).parent
-+
-+
- class Prismer(nn.Module):
-     def __init__(self, config):
-         super().__init__()
-@@ -26,7 +30,7 @@ class Prismer(nn.Module):
-             elif exp in ['obj_detection', 'ocr_detection']:
-                 self.experts[exp] = 64
- 
--        prismer_config = json.load(open('configs/prismer.json', 'r'))[config['prismer_model']]
-+        prismer_config = json.load(open(f'{cur_dir.parent}/configs/prismer.json', 'r'))[config['prismer_model']]
-         roberta_config = RobertaConfig.from_dict(prismer_config['roberta_model'])
- 
-         self.tokenizer = RobertaTokenizer.from_pretrained(prismer_config['roberta_model']['model_name'])
-@@ -35,7 +39,7 @@ class Prismer(nn.Module):
- 
-         self.prepare_to_train(config['freeze'])
-         self.ignored_modules = self.get_ignored_modules(config['freeze'])
--    
-+
-     def prepare_to_train(self, mode='none'):
-         for name, params in self.named_parameters():
-             if mode == 'freeze_lang':

prismer_model.py

@@ -79,28 +79,14 @@ class Model:
         if exp_name == self.exp_name:
             return
         config = {
-            'dataset':
-            'demo',
-            'data_path':
-            'prismer/helpers',
-            'label_path':
-            'prismer/helpers/labels',
-            'experts': [
-                'depth',
-                'normal',
-                'seg_coco',
-                'edge',
-                'obj_detection',
-                'ocr_detection',
-            ],
-            'image_resolution':
-            480,
-            'prismer_model':
-            'prismer_base',
-            'freeze':
-            'freeze_vision',
-            'prefix':
-            'A picture of',
+            'dataset': 'demo',
+            'data_path': 'prismer/helpers',
+            'label_path': 'prismer/helpers/labels',
+            'experts': ['depth', 'normal', 'seg_coco', 'edge', 'obj_detection', 'ocr_detection'],
+            'image_resolution': 480,
+            'prismer_model': 'prismer_base',
+            'freeze': 'freeze_vision',
+            'prefix': 'A picture of',
         }
         model = PrismerCaption(config)
         state_dict = torch.load(
@@ -118,27 +104,17 @@ class Model:
     @torch.inference_mode()
     def run_caption_model(self, exp_name: str) -> str:
         self.set_model(exp_name)
-
         _, test_dataset = create_dataset('caption', self.config)
-        test_loader = create_loader(test_dataset,
-                                    batch_size=1,
-                                    num_workers=4,
-                                    train=False)
+        test_loader = create_loader(test_dataset, batch_size=1, num_workers=4, train=False)
         experts, _ = next(iter(test_loader))
-        captions = self.model(experts,
-                              train=False,
-                              prefix=self.config['prefix'])
-        captions = self.tokenizer(captions,
-                                  max_length=30,
-                                  padding='max_length',
-                                  return_tensors='pt').input_ids
+        captions = self.model(experts, train=False, prefix=self.config['prefix'])
+        captions = self.tokenizer(captions, max_length=30, padding='max_length', return_tensors='pt').input_ids
         caption = captions.to(experts['rgb'].device)[0]
         caption = self.tokenizer.decode(caption, skip_special_tokens=True)
         caption = caption.capitalize() + '.'
         return caption
 
-    def run_caption(self, image_path: str,
-                    model_name: str) -> tuple[str | None, ...]:
+    def run_caption(self, image_path: str, model_name: str) -> tuple[str | None, ...]:
         out_paths = run_experts(image_path)
         caption = self.run_caption_model(model_name)
         return caption, *out_paths

style.css

@@ -1,3 +0,0 @@
-h1 {
-  text-align: center;
-}