detect-image / app.py
Mahmoud-abbasi-m's picture
Create app.py
09cc233 verified
import torch
from PIL import Image
from transformers import AutoModel, AutoTokenizer
class ImageChatbot:
def __init__(self, model_name='openbmb/MiniCPM-Llama3-V-2_5-int4'):
self.model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
self.model.eval()
def load_image(self, image_path):
image = Image.open(image_path).convert('RGB')
return image
def display_image(self, image):
import matplotlib.pyplot as plt
plt.imshow(image)
plt.axis('off')
plt.show()
def chat_with_image(self, image_path, question, sampling=True, temperature=0.7):
image = self.load_image(image_path)
self.display_image(image)
msgs = [{'role': 'user', 'content': question}]
res = self.model.chat(
image=image,
msgs=msgs,
tokenizer=self.tokenizer,
sampling=sampling,
temperature=temperature,
)
generated_text = ""
for new_text in res:
generated_text += new_text
print(new_text, flush=True, end='')
return generated_text
# Example usage
if __name__ == "__main__":
image_chatbot = ImageChatbot()
image_path = '/content/sample_data/Cat-profile-picture-35.jpg'
question = 'این شکل چی هست ؟'
generated_text = image_chatbot.chat_with_image(image_path, question)