Apoorva96nbd commited on
Commit
01259cd
·
verified ·
1 Parent(s): 9ba130a

added initial files

Browse files
Files changed (6) hide show
  1. .gitattributes +35 -35
  2. README.md +70 -13
  3. app.py +81 -0
  4. paligemma_tokenizer.model +3 -0
  5. requirements.txt +13 -0
  6. utils.py +204 -0
.gitattributes CHANGED
@@ -1,35 +1,35 @@
1
- *.7z filter=lfs diff=lfs merge=lfs -text
2
- *.arrow filter=lfs diff=lfs merge=lfs -text
3
- *.bin filter=lfs diff=lfs merge=lfs -text
4
- *.bz2 filter=lfs diff=lfs merge=lfs -text
5
- *.ckpt filter=lfs diff=lfs merge=lfs -text
6
- *.ftz filter=lfs diff=lfs merge=lfs -text
7
- *.gz filter=lfs diff=lfs merge=lfs -text
8
- *.h5 filter=lfs diff=lfs merge=lfs -text
9
- *.joblib filter=lfs diff=lfs merge=lfs -text
10
- *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
- *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
- *.model filter=lfs diff=lfs merge=lfs -text
13
- *.msgpack filter=lfs diff=lfs merge=lfs -text
14
- *.npy filter=lfs diff=lfs merge=lfs -text
15
- *.npz filter=lfs diff=lfs merge=lfs -text
16
- *.onnx filter=lfs diff=lfs merge=lfs -text
17
- *.ot filter=lfs diff=lfs merge=lfs -text
18
- *.parquet filter=lfs diff=lfs merge=lfs -text
19
- *.pb filter=lfs diff=lfs merge=lfs -text
20
- *.pickle filter=lfs diff=lfs merge=lfs -text
21
- *.pkl filter=lfs diff=lfs merge=lfs -text
22
- *.pt filter=lfs diff=lfs merge=lfs -text
23
- *.pth filter=lfs diff=lfs merge=lfs -text
24
- *.rar filter=lfs diff=lfs merge=lfs -text
25
- *.safetensors filter=lfs diff=lfs merge=lfs -text
26
- saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
- *.tar.* filter=lfs diff=lfs merge=lfs -text
28
- *.tar filter=lfs diff=lfs merge=lfs -text
29
- *.tflite filter=lfs diff=lfs merge=lfs -text
30
- *.tgz filter=lfs diff=lfs merge=lfs -text
31
- *.wasm filter=lfs diff=lfs merge=lfs -text
32
- *.xz filter=lfs diff=lfs merge=lfs -text
33
- *.zip filter=lfs diff=lfs merge=lfs -text
34
- *.zst filter=lfs diff=lfs merge=lfs -text
35
- *tfevents* filter=lfs diff=lfs merge=lfs -text
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
README.md CHANGED
@@ -1,13 +1,70 @@
1
- ---
2
- title: Vision
3
- emoji: 🌍
4
- colorFrom: indigo
5
- colorTo: yellow
6
- sdk: streamlit
7
- sdk_version: 1.37.1
8
- app_file: app.py
9
- pinned: false
10
- license: gemma
11
- ---
12
-
13
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ## Inhabitr Design Engine 1.0
2
+
3
+ #### PaliGemma Vision Model(unifying attributes model)
4
+ Leveraging the PaliGemma vision model to unify attribute extraction, this custom trained Model identifies and outputs furniture attributes—such as primary color and secondary color, primary material, secondary material, and design style—in JSON format.
5
+
6
+ # Features
7
+ Paligemma Model: A 3-billion parameter multi-model capable of generating captions from images, performing object detection, object segmentation ,OCR and VQA.
8
+ Custom Training: The model is custom trained on approximately 500 pairs of cropped images (sofa, accent chair) along with their captions.
9
+
10
+ # Requirements
11
+ Python 3.9+
12
+
13
+ Minimum 12GB of GPU for local inference
14
+
15
+ Pip (Python package installer)
16
+
17
+ # Installation
18
+ First, clone the repository:
19
+
20
+ git clone git clone https://Apoorva_inhabitr@bitbucket.org/Inhabitr/vision_models.git
21
+
22
+ # Create and Activate Python Virtual Environment
23
+
24
+ For Unix/macOS
25
+
26
+ python3 -m venv env
27
+ source env/bin/activate
28
+
29
+ For Windows
30
+
31
+ python -m venv env
32
+ .\env\Scripts\activate
33
+
34
+ # Install the required Python libraries
35
+
36
+ pip install -r requirements.txt
37
+
38
+ # Download Model Weights and Tokenizer
39
+
40
+ Download the model weights and tokenizer from Google Drive
41
+
42
+ https://drive.google.com/drive/folders/1S-z374V-yd3izeBatAMQlitmZ0mbMv6s?usp=drive_link.
43
+
44
+ Place the downloaded files in the models/ folder located in the root directory of the project.
45
+
46
+ # Start the API
47
+ python main.py
48
+ If running locally, the API will be available at http://localhost:5000.
49
+
50
+ # API Endpoints for Captions
51
+ predict attributes
52
+
53
+ POST /predict
54
+
55
+ Request Parameters
56
+
57
+ Send an image file
58
+
59
+ image: The image file used as the basis for the search.
60
+
61
+ Example Request with Image File
62
+
63
+ curl -X POST http://127.0.0.1:5000/predict \
64
+ -H "Content-Type: multipart/form-data" \
65
+ -F "file=@/path/to/your/image.jpg"
66
+
67
+ # Performance Considerations
68
+ GPU: Model parameters are loaded into memory for fast inference within 1 sec after initial loading.
69
+
70
+ CPU: Initial model loading may take 4-6 minutes based on your system configuration.
app.py ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from PIL import Image
3
+ import os
4
+ from fastapi import FastAPI, File, UploadFile
5
+ from pydantic import BaseModel
6
+ import uvicorn
7
+ import threading
8
+ from werkzeug.utils import secure_filename
9
+ from utils import setup_and_predict
10
+ from io import BytesIO
11
+
12
+ # Configure upload folder and allowed extensions
13
+ UPLOAD_FOLDER = 'uploads'
14
+ ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
15
+
16
+ if not os.path.exists(UPLOAD_FOLDER):
17
+ os.makedirs(UPLOAD_FOLDER)
18
+
19
+ def allowed_file(filename):
20
+ return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
21
+
22
+ def save_uploaded_file(uploaded_file):
23
+ filename = secure_filename(uploaded_file.name)
24
+ filepath = os.path.join(UPLOAD_FOLDER, filename)
25
+ with open(filepath, "wb") as f:
26
+ f.write(uploaded_file.getbuffer())
27
+ return filepath
28
+
29
+ # FastAPI app for Postman requests
30
+ api = FastAPI()
31
+
32
+ class PredictionResponse(BaseModel):
33
+ filename: str
34
+ result: str
35
+
36
+ @api.post("/predict", response_model=PredictionResponse)
37
+ async def predict_api(file: UploadFile = File(...)):
38
+ if allowed_file(file.filename):
39
+ filename = secure_filename(file.filename)
40
+ filepath = os.path.join(UPLOAD_FOLDER, filename)
41
+
42
+ with open(filepath, "wb") as f:
43
+ f.write(await file.read())
44
+
45
+ try:
46
+ result = setup_and_predict(filepath)
47
+ except Exception as e:
48
+ result = f"Unable to process the request: {e}"
49
+
50
+ return PredictionResponse(filename=filename, result=result)
51
+
52
+ return {"error": "Invalid file type"}
53
+
54
+ def run_api():
55
+ uvicorn.run(api, host="0.0.0.0", port=8000)
56
+
57
+ # Streamlit app
58
+ st.title("Upload an Image for Prediction")
59
+
60
+ uploaded_file = st.file_uploader("Choose an image...", type=['png', 'jpg', 'jpeg', 'gif'])
61
+
62
+ if uploaded_file is not None and allowed_file(uploaded_file.name):
63
+ filepath = save_uploaded_file(uploaded_file)
64
+
65
+ try:
66
+ result = setup_and_predict(filepath)
67
+ st.image(filepath, caption='Uploaded Image.', use_column_width=True)
68
+ st.write("Prediction Result:")
69
+ st.write(result)
70
+ except Exception as e:
71
+ st.error(f"Unable to process the request: {e}")
72
+
73
+ # Option to download the image
74
+ with open(filepath, "rb") as file:
75
+ st.download_button(label="Download Image", data=file, file_name=uploaded_file.name)
76
+
77
+ else:
78
+ st.warning("Please upload a valid image file.")
79
+
80
+ # Run FastAPI in a separate thread
81
+ threading.Thread(target=run_api).start()
paligemma_tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8986bb4f423f07f8c7f70d0dbe3526fb2316056c17bae71b1ea975e77a168fc6
3
+ size 4264023
requirements.txt ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Flask
2
+ jax
3
+ jaxlib
4
+ numpy
5
+ tensorflow
6
+ sentencepiece
7
+ Pillow
8
+ ml_collections
9
+ overrides
10
+ einops~=0.7
11
+ numpy>=1.26
12
+ flax
13
+ pillow
utils.py ADDED
@@ -0,0 +1,204 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import sys
3
+ import base64
4
+ import functools
5
+ import html
6
+ import io
7
+ import warnings
8
+ import jax
9
+ import jax.numpy as jnp
10
+ import numpy as np
11
+ import ml_collections
12
+ import tensorflow as tf
13
+ import sentencepiece
14
+ from PIL import Image
15
+
16
+ # TPUs with
17
+ if "COLAB_TPU_ADDR" in os.environ:
18
+ raise "It seems you are using Colab with remote TPUs which is not supported."
19
+ # Append big_vision code to python import path
20
+ if "big_vision_repo" not in sys.path:
21
+ sys.path.append("big_vision_repo")
22
+
23
+ # Import model definition from big_vision
24
+ from big_vision.models.proj.paligemma import paligemma
25
+ from big_vision.trainers.proj.paligemma import predict_fns
26
+ # Import big vision utilities
27
+ import big_vision.datasets.jsonl
28
+ import big_vision.utils
29
+ import big_vision.sharding
30
+ # Don't let TF use the GPU or TPUs
31
+ tf.config.set_visible_devices([], "GPU")
32
+ tf.config.set_visible_devices([], "TPU")
33
+ backend = jax.lib.xla_bridge.get_backend()
34
+
35
+ model_path = './Sofa-attributes-paligemma-ckpt.npz'
36
+ tokenizer_path = './paligemma_tokenizer.model'
37
+
38
+ # Define model
39
+ model_config = ml_collections.FrozenConfigDict({
40
+ "llm": {"vocab_size": 257_152},
41
+ "img": {"variant": "So400m/14", "pool_type": "none", "scan": True, "dtype_mm": "float16"}
42
+ })
43
+ model = paligemma.Model(**model_config)
44
+ tokenizer = sentencepiece.SentencePieceProcessor(tokenizer_path)
45
+ # Load params - this can take up to 1 minute in T4 colabs.
46
+ params = paligemma.load(None, model_path, model_config)
47
+ # Define `decode` function to sample outputs from the model.
48
+ decode_fn = predict_fns.get_all(model)['decode']
49
+ decode = functools.partial(decode_fn, devices=jax.devices(), eos_token=tokenizer.eos_id())
50
+
51
+ # Create a pytree mask of the trainable params.
52
+ def is_trainable_param(name, param): # pylint: disable=unused-argument
53
+ if name.startswith("llm/layers/attn/"): return True
54
+ if name.startswith("llm/"): return False
55
+ if name.startswith("img/"): return False
56
+ raise ValueError(f"Unexpected param name {name}")
57
+ trainable_mask = big_vision.utils.tree_map_with_names(is_trainable_param, params)
58
+ # If more than one device is available (e.g. multiple GPUs) the parameters can
59
+ # be sharded across them to reduce HBM usage per device.
60
+ mesh = jax.sharding.Mesh(jax.devices(), ("data"))
61
+
62
+ data_sharding = jax.sharding.NamedSharding(
63
+ mesh, jax.sharding.PartitionSpec("data"))
64
+
65
+ params_sharding = big_vision.sharding.infer_sharding(
66
+ params, strategy=[('.*', 'fsdp(axis="data")')], mesh=mesh)
67
+
68
+ # Yes: Some donated buffers are not usable.
69
+ warnings.filterwarnings(
70
+ "ignore", message="Some donated buffers were not usable")
71
+
72
+ @functools.partial(jax.jit, donate_argnums=(0,), static_argnums=(1,))
73
+ def maybe_cast_to_f32(params, trainable):
74
+ return jax.tree.map(lambda p, m: p.astype(jnp.float32) if m else p,
75
+ params, trainable)
76
+
77
+ # Loading all params in simultaneous - albeit much faster and more succinct -
78
+ # requires more RAM than the T4 colab runtimes have by default (12GB RAM).
79
+ # Instead we do it param by param.
80
+ params, treedef = jax.tree.flatten(params)
81
+ sharding_leaves = jax.tree.leaves(params_sharding)
82
+ trainable_leaves = jax.tree.leaves(trainable_mask)
83
+ for idx, (sharding, trainable) in enumerate(zip(sharding_leaves, trainable_leaves)):
84
+ params[idx] = big_vision.utils.reshard(params[idx], sharding)
85
+ params[idx] = maybe_cast_to_f32(params[idx], trainable)
86
+ params[idx].block_until_ready()
87
+ params = jax.tree.unflatten(treedef, params)
88
+
89
+ # Print params to show what the model is made of.
90
+ def parameter_overview(params):
91
+ for path, arr in big_vision.utils.tree_flatten_with_names(params)[0]:
92
+ print(f"{path:80s} {str(arr.shape):22s} {arr.dtype}")
93
+
94
+ print(" == Model params == ")
95
+ parameter_overview(params)
96
+ def setup_and_predict(image_path):
97
+ # Preprocess image and tokens
98
+ def preprocess_image(image, size=224):
99
+ # Model has been trained to handle images of different aspects ratios
100
+ # resized to 224x224 in the range [-1, 1]. Bilinear and antialias resize
101
+ # options are helpful to improve quality in some tasks.
102
+ image = np.asarray(image)
103
+ if image.ndim == 2: # Convert image without last channel into greyscale.
104
+ image = np.stack((image,)*3, axis=-1)
105
+ image = image[..., :3] # Remove alpha layer.
106
+ assert image.shape[-1] == 3
107
+
108
+ image = tf.constant(image)
109
+ image = tf.image.resize(image, (size, size), method='bilinear', antialias=True)
110
+ return image.numpy() / 127.5 - 1.0 # [0, 255]->[-1,1]
111
+
112
+ def preprocess_tokens(prefix, suffix=None, seqlen=None):
113
+ # Model has been trained to handle tokenized text composed of a prefix with
114
+ # full attention and a suffix with causal attention.
115
+ separator = "\n"
116
+ tokens = tokenizer.encode(prefix, add_bos=True) + tokenizer.encode(separator)
117
+ mask_ar = [0] * len(tokens) # 0 to use full attention for prefix.
118
+ mask_loss = [0] * len(tokens) # 0 to not use prefix tokens in the loss.
119
+
120
+ if suffix:
121
+ suffix = tokenizer.encode(suffix, add_eos=True)
122
+ tokens += suffix
123
+ mask_ar += [1] * len(suffix) # 1 to use causal attention for suffix.
124
+ mask_loss += [1] * len(suffix) # 1 to use suffix tokens in the loss.
125
+
126
+ mask_input = [1] * len(tokens) # 1 if its a token, 0 if padding.
127
+ if seqlen:
128
+ padding = [0] * max(0, seqlen - len(tokens))
129
+ tokens = tokens[:seqlen] + padding
130
+ mask_ar = mask_ar[:seqlen] + padding
131
+ mask_loss = mask_loss[:seqlen] + padding
132
+ mask_input = mask_input[:seqlen] + padding
133
+
134
+ return jax.tree.map(np.array, (tokens, mask_ar, mask_loss, mask_input))
135
+
136
+ def postprocess_tokens(tokens):
137
+ tokens = tokens.tolist() # np.array to list[int]
138
+ try: # Remove tokens at and after EOS if any.
139
+ eos_pos = tokens.index(tokenizer.eos_id())
140
+ tokens = tokens[:eos_pos]
141
+ except ValueError:
142
+ pass
143
+ return tokenizer.decode(tokens)
144
+
145
+ # Make predictions
146
+ # Evaluation/inference loop.
147
+ SEQLEN = 128
148
+ def make_predictions(data_iterator, *, num_examples=None,
149
+ batch_size=4, seqlen=SEQLEN, sampler="greedy"):
150
+ outputs = []
151
+ while True:
152
+ # Construct a list of examples in the batch.
153
+ examples = []
154
+ try:
155
+ for _ in range(batch_size):
156
+ examples.append(next(data_iterator))
157
+ examples[-1]["_mask"] = np.array(True) # Indicates true example.
158
+ except StopIteration:
159
+ if len(examples) == 0:
160
+ return outputs
161
+
162
+ # Not enough examples to complete a batch. Pad by repeating last example.
163
+ while len(examples) % batch_size:
164
+ examples.append(dict(examples[-1]))
165
+ examples[-1]["_mask"] = np.array(False) # Indicates padding example.
166
+
167
+ # Convert list of examples into a dict of np.arrays and load onto devices.
168
+ batch = jax.tree.map(lambda *x: np.stack(x), *examples)
169
+ batch = big_vision.utils.reshard(batch, data_sharding)
170
+
171
+ # Make model predictions
172
+ tokens = decode({"params": params}, batch=batch,
173
+ max_decode_len=seqlen, sampler=sampler)
174
+
175
+ # Fetch model predictions to device and detokenize.
176
+ tokens, mask = jax.device_get((tokens, batch["_mask"]))
177
+ tokens = tokens[mask] # remove padding examples.
178
+ responses = [postprocess_tokens(t) for t in tokens]
179
+
180
+ # Append to html output.
181
+ for example, response in zip(examples, responses):
182
+ outputs.append((example["image"], response))
183
+ if num_examples and len(outputs) >= num_examples:
184
+ return outputs
185
+
186
+ def test_data_iterator(file_name):
187
+ image = Image.open(file_name)
188
+ image = preprocess_image(image)
189
+
190
+ prefix = "caption en"
191
+ tokens, mask_ar, _, mask_input = preprocess_tokens(prefix, seqlen=SEQLEN)
192
+
193
+ yield {
194
+ "image": np.asarray(image),
195
+ "text": np.asarray(tokens),
196
+ "mask_ar": np.asarray(mask_ar),
197
+ "mask_input": np.asarray(mask_input)
198
+ }
199
+
200
+ # Call the prediction function and print the result
201
+ image, caption = make_predictions(test_data_iterator(file_name=image_path), batch_size=1)[0]
202
+ return caption
203
+
204
+