Spaces:
Running
Running
File size: 3,248 Bytes
ba395cf |
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 |
import base64
import requests
from io import BytesIO
from PIL import Image
import gradio as gr
from pymongo import MongoClient
import time
import uuid
# MongoDB setup
mongo_client = MongoClient("mongodb+srv://atharva2021:[email protected]/")
db = mongo_client['bajaj']
collection = db['client']
# Function to encode the image to base64
def encode_image(img):
buffered = BytesIO()
img.save(buffered, format="PNG")
encoded_string = base64.b64encode(buffered.getvalue()).decode("utf-8")
return encoded_string
# Chat function with Pixtral and MongoDB saving
def chat_with_pixtral(uploaded_file, mrn_number, user_question):
if uploaded_file is not None and mrn_number.strip() != "":
base64_img = encode_image(uploaded_file)
api = "https://api.hyperbolic.xyz/v1/chat/completions"
api_key = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJyZzMyNzAyNEBnbWFpbC5jb20ifQ._frFve-BYZdb0Qo6FIj6xcDcxpY-6QlC2O-ToQxBjkc" # Add your API key here
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}",
}
payload = {
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": user_question},
{
"type": "image_url",
"image_url": {"url": f"data:image/jpeg;base64,{base64_img}"},
},
],
}
],
"model": "mistralai/Pixtral-12B-2409",
"max_tokens": 2048,
"temperature": 0.7,
"top_p": 0.9,
}
response = requests.post(api, headers=headers, json=payload)
# Process API response
if response.status_code == 200:
response_data = response.json()
if 'choices' in response_data:
assistant_response = response_data['choices'][0]['message']['content']
else:
assistant_response = "Response format is incorrect"
else:
assistant_response = f"API request failed: {response.status_code} - {response.text}"
# Generate a unique ID for the request
unique_id = str(uuid.uuid4())
# Save the result to MongoDB with the specified format
document = {
'mrn_number': mrn_number,
'ocr_result': assistant_response, # This will be the OCR/API result
'unique_id': unique_id,
'got_mode': "plain texts OCR",
'timestamp': time.time()
}
collection.insert_one(document) # Insert the document into MongoDB
return assistant_response
return "Upload your image, enter MRN number, and enter your question."
# Gradio GUI
iface = gr.Interface(
fn=chat_with_pixtral,
inputs=[
gr.Image(type="pil", label="Upload Your Image"),
gr.Textbox(label="Enter MRN Number"),
gr.Textbox(label="Please enter your question")
],
outputs="text",
title="Pixtral Image Chat",
description="Upload your Image, enter MRN number, and get insights out of the Image"
)
iface.launch(share=True)
|