File size: 9,398 Bytes
39b01b9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1072512
39b01b9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
20ceca0
39b01b9
20ceca0
39b01b9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5b02197
 
39b01b9
 
 
 
 
 
 
 
 
 
20ceca0
39b01b9
 
 
 
 
 
 
 
20ceca0
39b01b9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
import os
import torch
import torch.nn as nn
import pandas as pd
from PIL import Image
from torchvision import transforms
from transformers import BertTokenizer, AutoModel
from torch.utils.data import Dataset, DataLoader, random_split
from sklearn.model_selection import train_test_split
from typing import List
from dataclasses import dataclass
import gradio as gr
import torch, re
import numpy as np
from transformers import WhisperProcessor, WhisperForConditionalGeneration, ViTImageProcessor, BertTokenizer, BlipProcessor, BlipForQuestionAnswering, AutoProcessor, AutoModelForCausalLM, DonutProcessor, VisionEncoderDecoderModel, Pix2StructProcessor, Pix2StructForConditionalGeneration, AutoModelForSeq2SeqLM

import librosa
from PIL import Image
from torch.nn.utils import rnn
from gtts import gTTS

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class LabelClassifier(nn.Module):
    def __init__(self):
        super(LabelClassifier, self).__init__()
        self.text_encoder = AutoModel.from_pretrained('bert-base-uncased')
        self.image_encoder = AutoModel.from_pretrained('microsoft/swin-tiny-patch4-window7-224')
        self.intermediate_dim = 128
        self.fusion = nn.Sequential(
            nn.Linear(self.text_encoder.config.hidden_size + self.image_encoder.config.hidden_size, self.intermediate_dim),
            nn.ReLU(),
            nn.Dropout(0.5),
        )
        self.classifier = nn.Linear(self.intermediate_dim, 6)  # Concatenating BERT output and Swin Transformer output

        self.criterion = nn.CrossEntropyLoss()


    def forward(self,
        input_ids: torch.LongTensor,pixel_values: torch.FloatTensor, attention_mask: torch.LongTensor = None, token_type_ids: torch.LongTensor = None, labels: torch.LongTensor = None):

        encoded_text = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
        encoded_image = self.image_encoder(pixel_values=pixel_values)

        # print(encoded_text['last_hidden_state'].shape)
        # print(encoded_image['last_hidden_state'].shape)

        fused_state = self.fusion(torch.cat((encoded_text['pooler_output'], encoded_image['pooler_output']), dim=1))


        # Pass through the classifier
        logits = self.classifier(fused_state)

        out = {"logits": logits}

        if labels is not None:
            loss = self.criterion(logits, labels)
            out["loss"] = loss


        return out

model = LabelClassifier().to(device)
model.load_state_dict(torch.load('classifier.pth', map_location=torch.device('cpu')))



tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
processor = ViTImageProcessor.from_pretrained('microsoft/swin-tiny-patch4-window7-224')


# Load the Whisper model in Hugging Face format:
# processor2 = WhisperProcessor.from_pretrained("openai/whisper-medium.en")
# model2 = WhisperForConditionalGeneration.from_pretrained("openai/whisper-medium.en")



def m1(que, image):
    processor3 = BlipProcessor.from_pretrained("Salesforce/blip-vqa-capfilt-large")
    model3 = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-capfilt-large")

    inputs = processor3(image, que, return_tensors="pt")

    out = model3.generate(**inputs)
    return processor3.decode(out[0], skip_special_tokens=True)

def m2(que, image):
    processor3 = AutoProcessor.from_pretrained("microsoft/git-large-textvqa")
    model3 = AutoModelForCausalLM.from_pretrained("microsoft/git-large-textvqa")

    pixel_values = processor3(images=image, return_tensors="pt").pixel_values

    input_ids = processor3(text=que, add_special_tokens=False).input_ids
    input_ids = [processor3.tokenizer.cls_token_id] + input_ids
    input_ids = torch.tensor(input_ids).unsqueeze(0)

    generated_ids = model3.generate(pixel_values=pixel_values, input_ids=input_ids, max_length=50)
    return processor3.batch_decode(generated_ids, skip_special_tokens=True)

def m3(que, image):
    processor3 = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa")
    model3 = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa")

    model3.to(device)

    prompt = "<s_docvqa><s_question>{que}</s_question><s_answer>"
    decoder_input_ids = processor3.tokenizer(prompt, add_special_tokens=False, return_tensors="pt").input_ids

    pixel_values = processor3(image, return_tensors="pt").pixel_values

    outputs = model3.generate(
        pixel_values.to(device),
        decoder_input_ids=decoder_input_ids.to(device),
        max_length=model3.decoder.config.max_position_embeddings,
        pad_token_id=processor3.tokenizer.pad_token_id,
        eos_token_id=processor3.tokenizer.eos_token_id,
        use_cache=True,
        bad_words_ids=[[processor3.tokenizer.unk_token_id]],
        return_dict_in_generate=True,
    )

    sequence = processor3.batch_decode(outputs.sequences)[0]
    sequence = sequence.replace(processor3.tokenizer.eos_token, "").replace(processor3.tokenizer.pad_token, "")
    sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # remove first task start token
    return processor3.token2json(sequence)['answer']

def m4(que, image):
    processor3 = Pix2StructProcessor.from_pretrained('google/matcha-plotqa-v2')
    model3 = Pix2StructForConditionalGeneration.from_pretrained('google/matcha-plotqa-v2')

    inputs = processor3(images=image, text=que, return_tensors="pt")
    predictions = model3.generate(**inputs, max_new_tokens=512)
    return processor3.decode(predictions[0], skip_special_tokens=True)

def m5(que, image):

    processor3 = AutoProcessor.from_pretrained("google/pix2struct-ocrvqa-large")
    model3 = AutoModelForSeq2SeqLM.from_pretrained("google/pix2struct-ocrvqa-large")

    inputs = processor3(images=image, text=que, return_tensors="pt")

    predictions = model3.generate(**inputs)
    return processor3.decode(predictions[0], skip_special_tokens=True)

def m6(que, image):
    processor3 = AutoProcessor.from_pretrained("google/pix2struct-infographics-vqa-large")
    model3 = AutoModelForSeq2SeqLM.from_pretrained("google/pix2struct-infographics-vqa-large")

    inputs = processor3(images=image, text=que, return_tensors="pt")

    predictions = model3.generate(**inputs)
    return processor3.decode(predictions[0], skip_special_tokens=True)


def predict_answer(category, que, image):
    if category == 0:
        return m1(que, image)
    elif category == 1:
        return m2(que, image)
    elif category == 2:
        return m3(que, image)
    elif category == 3:
        return m4(que, image)
    elif category == 4:
        return m5(que, image)
    else:
        return m6(que, image)



def transcribe_audio(audio):
    # print(audio)
    processor2 = WhisperProcessor.from_pretrained("openai/whisper-large-v3",language='en')
    model2 = WhisperForConditionalGeneration.from_pretrained("openai/whisper-large-v3")

    sampling_rate = audio[0]
    audio_data = audio[1]

    # print(np.array([audio_data]).shape)
    audio_data_float = np.array(audio_data).astype(np.float32)
    resampled_audio_data = librosa.resample(audio_data_float, orig_sr=sampling_rate, target_sr=16000)


    # Use the model and processor to transcribe the audio:
    input_features = processor2(
        resampled_audio_data, sampling_rate=16000, return_tensors="pt"
    ).input_features

    # Generate token ids
    predicted_ids = model2.generate(input_features)

    # Decode token ids to text
    transcription = processor2.batch_decode(predicted_ids, skip_special_tokens=True)[0]

    return transcription


def predict_category(que, input_image):
    # print(type(input_image))
    # print(input_image)

    encoded_text = tokenizer(
        text=que,
        padding='longest',
        max_length=24,
        truncation=True,
        return_tensors='pt',
        return_token_type_ids=True,
        return_attention_mask=True,
    )

    encoded_image = processor(input_image, return_tensors='pt').to(device)

    dict = {
        'input_ids':  encoded_text['input_ids'].to(device),
        'token_type_ids': encoded_text['token_type_ids'].to(device),
        'attention_mask': encoded_text['attention_mask'].to(device),
        'pixel_values': encoded_image['pixel_values'].to(device)
    }

    output = model(input_ids=dict['input_ids'],token_type_ids=dict['token_type_ids'],attention_mask=dict['attention_mask'],pixel_values=dict['pixel_values'])

    preds = output["logits"].argmax(axis=-1).cpu().numpy()

    return preds[0]


def combine(audio, input_image):
    que = transcribe_audio(audio)
    # que = "What is the animal here?"

    image = Image.fromarray(input_image).convert('RGB')
    category = predict_category(que, image)

    answer = predict_answer(0, que, image)

    # print(category)

    tts = gTTS(answer)
    tts.save('answer.mp3')
    return que, answer, 'answer.mp3'



# Define the Gradio interface for recording audio and displaying the transcription
model_interface = gr.Interface(fn=combine, inputs=[gr.Microphone(label="Ask your question"),gr.Image(label="Upload the image")], outputs=[gr.Text(label="Transcribed Question"), gr.Text(label="Answer"), gr.Audio(label="Audio Answer")])
# image_upload_interface = gr.Interface(fn=upload_image, inputs=gr.Image(label="Upload the image"), outputs="text")

# Launch the Gradio interface
model_interface.launch(debug=True)