saurabhharak commited on
Commit
5deb389
Β·
1 Parent(s): e9b8ff0

code added

Browse files
Files changed (4) hide show
  1. .gitignore +162 -0
  2. model_init.py +44 -0
  3. model_use.py +62 -0
  4. streamilt_app.py +71 -0
.gitignore ADDED
@@ -0,0 +1,162 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[cod]
4
+ *$py.class
5
+
6
+ # C extensions
7
+ *.so
8
+
9
+ content/
10
+
11
+ # Distribution / packaging
12
+ .Python
13
+ build/
14
+ develop-eggs/
15
+ dist/
16
+ downloads/
17
+ eggs/
18
+ .eggs/
19
+ lib/
20
+ lib64/
21
+ parts/
22
+ sdist/
23
+ var/
24
+ wheels/
25
+ share/python-wheels/
26
+ *.egg-info/
27
+ .installed.cfg
28
+ *.egg
29
+ MANIFEST
30
+
31
+ # PyInstaller
32
+ # Usually these files are written by a python script from a template
33
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
34
+ *.manifest
35
+ *.spec
36
+
37
+ # Installer logs
38
+ pip-log.txt
39
+ pip-delete-this-directory.txt
40
+
41
+ # Unit test / coverage reports
42
+ htmlcov/
43
+ .tox/
44
+ .nox/
45
+ .coverage
46
+ .coverage.*
47
+ .cache
48
+ nosetests.xml
49
+ coverage.xml
50
+ *.cover
51
+ *.py,cover
52
+ .hypothesis/
53
+ .pytest_cache/
54
+ cover/
55
+
56
+ # Translations
57
+ *.mo
58
+ *.pot
59
+
60
+ # Django stuff:
61
+ *.log
62
+ local_settings.py
63
+ db.sqlite3
64
+ db.sqlite3-journal
65
+
66
+ # Flask stuff:
67
+ instance/
68
+ .webassets-cache
69
+
70
+ # Scrapy stuff:
71
+ .scrapy
72
+
73
+ # Sphinx documentation
74
+ docs/_build/
75
+
76
+ # PyBuilder
77
+ .pybuilder/
78
+ target/
79
+
80
+ # Jupyter Notebook
81
+ .ipynb_checkpoints
82
+
83
+ # IPython
84
+ profile_default/
85
+ ipython_config.py
86
+
87
+ # pyenv
88
+ # For a library or package, you might want to ignore these files since the code is
89
+ # intended to run in multiple environments; otherwise, check them in:
90
+ # .python-version
91
+
92
+ # pipenv
93
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
94
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
95
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
96
+ # install all needed dependencies.
97
+ #Pipfile.lock
98
+
99
+ # poetry
100
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
101
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
102
+ # commonly ignored for libraries.
103
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
104
+ #poetry.lock
105
+
106
+ # pdm
107
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
108
+ #pdm.lock
109
+ # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
110
+ # in version control.
111
+ # https://pdm.fming.dev/#use-with-ide
112
+ .pdm.toml
113
+
114
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
115
+ __pypackages__/
116
+
117
+ # Celery stuff
118
+ celerybeat-schedule
119
+ celerybeat.pid
120
+
121
+ # SageMath parsed files
122
+ *.sage.py
123
+
124
+ # Environments
125
+ .env
126
+ .venv
127
+ env/
128
+ venv/
129
+ ENV/
130
+ env.bak/
131
+ venv.bak/
132
+
133
+ # Spyder project settings
134
+ .spyderproject
135
+ .spyproject
136
+
137
+ # Rope project settings
138
+ .ropeproject
139
+
140
+ # mkdocs documentation
141
+ /site
142
+
143
+ # mypy
144
+ .mypy_cache/
145
+ .dmypy.json
146
+ dmypy.json
147
+
148
+ # Pyre type checker
149
+ .pyre/
150
+
151
+ # pytype static type analyzer
152
+ .pytype/
153
+
154
+ # Cython debug symbols
155
+ cython_debug/
156
+
157
+ # PyCharm
158
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
159
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
160
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
161
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
162
+ #.idea/
model_init.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import VisionEncoderDecoderModel,ViTFeatureExtractor,AutoTokenizer,ViTImageProcessor
2
+ import torch
3
+ from PIL import Image
4
+
5
+ model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
6
+ feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
7
+ tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
8
+
9
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
10
+ model.to(device)
11
+
12
+ max_length = 16
13
+ num_beams = 4
14
+ gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
15
+
16
+
17
+ def predict_step(image_paths):
18
+ images = []
19
+ for image_path in image_paths:
20
+ i_image = Image.open(image_path)
21
+ if i_image.mode != "RGB":
22
+ i_image = i_image.convert(mode="RGB")
23
+ images.append(i_image)
24
+ pixel_values = feature_extractor(images=images, return_tensors="pt").pixel_values
25
+ pixel_values = pixel_values.to(device)
26
+ output_ids = model.generate(pixel_values, **gen_kwargs)
27
+ preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
28
+ preds = [pred.strip() for pred in preds]
29
+ return preds
30
+
31
+
32
+ # Define paths to save the components in your Google Drive
33
+ drive_folder = "/content/drive/My Drive/image_captioning_streamlit"
34
+
35
+ saved_model_directory = f"{drive_folder}/saved_model"
36
+ saved_feature_extractor_directory = f"{drive_folder}/saved_feature_extractor"
37
+ saved_tokenizer_directory = f"{drive_folder}/saved_tokenizer"
38
+
39
+
40
+ # Save the model and its components
41
+ model.save_pretrained(saved_model_directory)
42
+ feature_extractor.save_pretrained(saved_feature_extractor_directory)
43
+ tokenizer.save_pretrained(saved_tokenizer_directory)
44
+
model_use.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
2
+ import torch
3
+ from PIL import Image
4
+
5
+
6
+ drive_folder = "/content/drive/My Drive/image_captioning_streamlit"
7
+
8
+ saved_model_directory = f"{drive_folder}/saved_model"
9
+ saved_feature_extractor_directory = f"{drive_folder}/saved_feature_extractor"
10
+ saved_tokenizer_directory = f"{drive_folder}/saved_tokenizer"
11
+ # Define paths to save the components in your Google Drive
12
+
13
+
14
+ saved_model = VisionEncoderDecoderModel.from_pretrained(saved_model_directory)
15
+ saved_feature_extractor = ViTImageProcessor.from_pretrained(saved_feature_extractor_directory)
16
+ saved_tokenizer = AutoTokenizer.from_pretrained(saved_tokenizer_directory)
17
+
18
+ # Move the model to the appropriate device (GPU if available)
19
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
20
+ saved_model.to(device)
21
+
22
+ # Define prediction parameters
23
+ max_length = 16
24
+ num_beams = 4
25
+ gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
26
+
27
+ # Define the prediction function
28
+ def predict_step(image_paths):
29
+ """
30
+ Generate predictions for a list of image paths.
31
+
32
+ Args:
33
+ image_paths (List[str]): A list of file paths to the images.
34
+
35
+ Returns:
36
+ List[str]: A list of predicted strings.
37
+
38
+ Raises:
39
+ None
40
+
41
+ Examples:
42
+ >>> image_paths = ["path/to/image1.jpg", "path/to/image2.jpg"]
43
+ >>> predict_step(image_paths)
44
+ ["prediction1", "prediction2"]
45
+ """
46
+ images = []
47
+ for image_path in image_paths:
48
+ i_image = Image.open(image_path)
49
+ if i_image.mode != "RGB":
50
+ i_image = i_image.convert(mode="RGB")
51
+ images.append(i_image)
52
+
53
+ pixel_values = saved_feature_extractor(images=images, return_tensors="pt").pixel_values
54
+ pixel_values = pixel_values.to(device)
55
+
56
+ output_ids = saved_model.generate(pixel_values, **gen_kwargs)
57
+
58
+ preds = saved_tokenizer.batch_decode(output_ids, skip_special_tokens=True)
59
+ preds = [pred.strip() for pred in preds]
60
+ return preds
61
+
62
+
streamilt_app.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
3
+ import torch
4
+ from PIL import Image
5
+
6
+ # Load the saved model and its components
7
+ saved_model_directory = "content/saved_model"
8
+ saved_feature_extractor_directory = "content/saved_feature_extractor"
9
+ saved_tokenizer_directory = "content/saved_tokenizer"
10
+
11
+ saved_model = VisionEncoderDecoderModel.from_pretrained(saved_model_directory)
12
+ saved_feature_extractor = ViTImageProcessor.from_pretrained(saved_feature_extractor_directory)
13
+ saved_tokenizer = AutoTokenizer.from_pretrained(saved_tokenizer_directory)
14
+
15
+ # Move the model to the appropriate device (GPU if available)
16
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
17
+ saved_model.to(device)
18
+
19
+ # Define prediction parameters
20
+ max_length = 16
21
+ num_beams = 4
22
+ gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
23
+
24
+ # Define the prediction function
25
+ def predict_step(image):
26
+ i_image = Image.open(image)
27
+ if i_image.mode != "RGB":
28
+ i_image = i_image.convert(mode="RGB")
29
+
30
+ pixel_values = saved_feature_extractor(images=[i_image], return_tensors="pt").pixel_values
31
+ pixel_values = pixel_values.to(device)
32
+
33
+ output_ids = saved_model.generate(pixel_values, **gen_kwargs)
34
+
35
+ preds = saved_tokenizer.batch_decode(output_ids, skip_special_tokens=True)
36
+ preds = [pred.strip() for pred in preds]
37
+ return preds[0]
38
+
39
+
40
+ def main():
41
+ # Streamlit app
42
+ st.set_page_config(
43
+ page_title="⭐ Image Captioning App",
44
+ page_icon="⭐",
45
+ layout="centered"
46
+ )
47
+
48
+ st.title("πŸ“· Image Captioning Using Transformers")
49
+ image = None
50
+ # Upload image or provide URL
51
+ st.write("Upload an image πŸ“€ or provide an image URL πŸ”—:")
52
+ uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
53
+ image_url = st.text_input("Or provide an image URL:")
54
+
55
+ if uploaded_image is not None:
56
+ image = uploaded_image
57
+ elif image_url:
58
+ image = image_url
59
+ else:
60
+ st.warning("❗ Please upload an image or provide an image URL.")
61
+
62
+ if image:
63
+ try:
64
+ caption = predict_step(image)
65
+ st.image(image, caption=f"πŸ“ Predicted caption: {caption}", use_column_width=True)
66
+ except Exception as e:
67
+ st.error("❌ An error occurred while generating the caption.")
68
+ st.error(str(e))
69
+
70
+ if __name__ == "__main__":
71
+ main()