ask-me-anything / app.py
kendrickfff's picture
Update app.py
5cbb24b verified
import os
import json
import gradio as gr
from ultralytics import YOLO
from PIL import Image, ImageDraw
import torch
from langchain_google_genai.chat_models import ChatGoogleGenerativeAI
# Load model YOLOv8
model = YOLO("yolov8n.pt")
# Load credentials (stringified JSON) from environment variable for Gemini
credentials_string = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
if not credentials_string:
raise ValueError("GOOGLE_APPLICATION_CREDENTIALS is not set in the environment!")
# Parse the stringified JSON back to a Python dictionary
credentials = json.loads(credentials_string)
# Save the credentials to a temporary JSON file (required by Google SDKs)
with open("service_account.json", "w") as f:
json.dump(credentials, f)
# Set the GOOGLE_APPLICATION_CREDENTIALS environment variable to the temporary file
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "service_account.json"
# Initialize Gemini model (chatbot)
llm = ChatGoogleGenerativeAI(model='gemini-1.5-pro')
# Global chat history
chat_history = []
# Fungsi untuk chatting dengan chatbot
def chat_with_bot(message):
global chat_history
response = llm.predict(message) # Menggunakan Gemini untuk menghasilkan respon
bot_response = f"Bot: {response}"
chat_history.append((message, bot_response))
return chat_history
# Fungsi untuk menganalisis gambar
def analyze_image(image_path):
global chat_history
try:
# Load gambar
image = Image.open(image_path).convert("RGB")
# Prediksi objek dalam gambar
results = model(image)
# Ambil hasil deteksi
detected_objects = []
image_draw = image.copy()
draw = ImageDraw.Draw(image_draw)
for result in results:
for box in result.boxes.data:
x1, y1, x2, y2, score, class_id = box.tolist()
if score > 0.5: # Hanya tampilkan objek dengan confidence score > 0.5
class_name = model.names[int(class_id)]
detected_objects.append(f"{class_name} (score: {score:.2f})")
draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
draw.text((x1, y1), class_name, fill="red")
if detected_objects:
bot_response = f"Objects detected: {', '.join(detected_objects)}."
else:
bot_response = "No objects detected."
chat_history.append(("Uploaded an image for analysis", bot_response))
return image_draw, chat_history
except Exception as e:
error_msg = f"Error processing the image: {str(e)}"
chat_history.append(("Error during image analysis", error_msg))
return None, chat_history
# Bangun antarmuka Gradio
with gr.Blocks() as demo:
gr.Markdown("# Ken Chatbot")
gr.Markdown("Ask me anything or upload an image for analysis!")
chatbot = gr.Chatbot(elem_id="chatbot")
msg = gr.Textbox(label="Type your message here...", placeholder="Enter your message...", show_label=False)
send_btn = gr.Button("Send")
img_upload = gr.Image(type="filepath", label="Upload an image for analysis")
img_output = gr.Image(label="Detected Objects")
msg.submit(chat_with_bot, msg, chatbot)
send_btn.click(chat_with_bot, msg, chatbot)
send_btn.click(lambda: "", None, msg) # Clear input field
img_upload.change(analyze_image, img_upload, [img_output, chatbot])
demo.launch()