Mahmoud-abbasi-m commited on
Commit
09cc233
1 Parent(s): 76bbb49

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +43 -0
app.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from PIL import Image
3
+ from transformers import AutoModel, AutoTokenizer
4
+
5
+ class ImageChatbot:
6
+ def __init__(self, model_name='openbmb/MiniCPM-Llama3-V-2_5-int4'):
7
+ self.model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
8
+ self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
9
+ self.model.eval()
10
+
11
+ def load_image(self, image_path):
12
+ image = Image.open(image_path).convert('RGB')
13
+ return image
14
+
15
+ def display_image(self, image):
16
+ import matplotlib.pyplot as plt
17
+ plt.imshow(image)
18
+ plt.axis('off')
19
+ plt.show()
20
+
21
+ def chat_with_image(self, image_path, question, sampling=True, temperature=0.7):
22
+ image = self.load_image(image_path)
23
+ self.display_image(image)
24
+ msgs = [{'role': 'user', 'content': question}]
25
+ res = self.model.chat(
26
+ image=image,
27
+ msgs=msgs,
28
+ tokenizer=self.tokenizer,
29
+ sampling=sampling,
30
+ temperature=temperature,
31
+ )
32
+ generated_text = ""
33
+ for new_text in res:
34
+ generated_text += new_text
35
+ print(new_text, flush=True, end='')
36
+ return generated_text
37
+
38
+ # Example usage
39
+ if __name__ == "__main__":
40
+ image_chatbot = ImageChatbot()
41
+ image_path = '/content/sample_data/Cat-profile-picture-35.jpg'
42
+ question = 'این شکل چی هست ؟'
43
+ generated_text = image_chatbot.chat_with_image(image_path, question)