CEEW / app.py
ceew36's picture
Update app.py
5c73743 verified
import PIL.Image as Image
import gradio as gr
from ultralytics import YOLO
import os
import time
import uuid
from langchain_groq import ChatGroq
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains import create_retrieval_chain
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyPDFLoader
from dotenv import load_dotenv
load_dotenv()
groq_api_key = os.getenv('GROQ_API_KEY')
# Initialize object detection model
model = YOLO("version4c.pt")
# Set default confidence and IoU thresholds
CONF_THRESHOLD = 0.25
IOU_THRESHOLD = 0.45
def predict_image(img):
# Perform object detection
results = model.predict(source=img, conf=CONF_THRESHOLD, iou=IOU_THRESHOLD, show_labels=True, show_conf=True, imgsz=640)
# Check if any objects were detected
if len(results[0].boxes) == 0:
return None, "Please upload a clearer image, and don't upload images of breeds that are not been used."
# Plot the result
for r in results:
im_array = r.plot()
im = Image.fromarray(im_array[..., ::-1])
# Generate a unique filename
filename = f"detected_result_{uuid.uuid4()}.jpg"
# Save the image as JPG
im.save(filename, format='JPEG')
return filename, None
# Initialize chatbot components
llm = ChatGroq(groq_api_key=groq_api_key, model_name="Llama3-70b-8192")
prompt = ChatPromptTemplate.from_template(
""" Answer the questions based on the provided context only. Please provide the most accurate response based on the question <context> {context} </context> Questions:{input} """
)
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
loader = PyPDFLoader("Document.pdf")
docs = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
final_documents = text_splitter.split_documents(docs)
# Extract text content from the Document instances
doc_texts = [doc.page_content for doc in final_documents]
embeddings_result = embeddings.embed_documents(doc_texts)
if embeddings_result:
vectors = FAISS.from_documents(final_documents, embeddings)
else:
raise ValueError("Failed to generate embeddings. Please check your input documents or try a different embedding model.")
document_chain = create_stuff_documents_chain(llm, prompt)
retriever = vectors.as_retriever()
retrieval_chain = create_retrieval_chain(retriever, document_chain)
def print_like_dislike(x: gr.LikeData):
print(x.index, x.value, x.liked)
def add_message(history, message):
if message is not None:
history.append((message, None))
return history, gr.Textbox(value=None, interactive=False)
stop_generation = False
def bot(history):
global stop_generation
stop_generation = False
message = history[-1][0]
start_time = time.time()
response = retrieval_chain.invoke({'input': message})['answer']
response_time = time.time() - start_time
if response_time > 6:
return [(f"Sorry, I couldn't generate a response within 6 seconds. Please try again.", None)]
history[-1][1] = ""
for character in response:
if stop_generation:
break
history[-1][1] += character
time.sleep(0.05)
yield history
def stop_response(dummy_placeholder):
global stop_generation
stop_generation = True
with gr.Blocks() as demo:
with gr.Row():
with gr.Column(scale=2):
model_input = gr.Image(type="pil", label="Upload Image")
model_output = gr.Image(type="filepath", label="Result")
caution_message = gr.Textbox(label="Caution", visible=False)
def process_image(img):
result, caution = predict_image(img)
if caution:
return None, caution, gr.Image(visible=False), gr.Textbox(visible=True)
else:
return result, "", gr.Image(visible=True), gr.Textbox(visible=False)
model_btn = gr.Button("Detect Result")
model_btn.click(process_image, inputs=model_input, outputs=[model_output, caution_message, model_output, caution_message])
with gr.Column(scale=1):
chatbot = gr.Chatbot(
[],
elem_id="chatbot",
bubble_full_width=False
)
chat_input = gr.Textbox(interactive=True, placeholder="Enter message...", show_label=False)
chat_msg = chat_input.submit(add_message, [chatbot, chat_input], [chatbot, chat_input])
bot_msg = chat_msg.then(bot, chatbot, chatbot, api_name="bot_response")
bot_msg.then(lambda: gr.Textbox(interactive=True), None, [chat_input])
chatbot.like(print_like_dislike, None, None)
stop_btn = gr.Button("Stop Generation")
stop_btn.click(stop_response, None, None)
demo.queue()
if __name__ == "__main__":
demo.launch()