collarvision / app.py
Kazel's picture
Upload app.py
df97b0c verified
raw
history blame
1.93 kB
import gradio as gr
import cv2
import ollama
import threading
# Initialize the webcam
cap = cv2.VideoCapture(0)
def query_the_image(query: str, image_list: list[str]):
try:
res = ollama.chat(
model='llava',
options={
'temperature': 0,
"top_k": 1,
'top_p': 0.1,
'mirostat_tau': 1.0,
'num_ctx': 1024,
'seed': 42,
'num_predict': 128
},
messages=[
{
'role': 'system',
'content': "You are a home surveillance system. Answer with very short sentences."
},
{
'role': 'user',
'content': query,
'images': image_list,
}
]
)
return res['message']['content']
except Exception as e:
return f"Error: {e}"
def get_frame():
ret, frame = cap.read()
if not ret:
return None
_, buffer = cv2.imencode('.jpg', frame)
return buffer.tobytes()
def process_image(prompt):
frame_data = get_frame()
if frame_data:
return query_the_image(prompt, [frame_data])
return "Error capturing image"
def video_feed():
while True:
ret, frame = cap.read()
if ret:
yield cv2.imencode('.jpg', frame)[1].tobytes()
else:
break
gui = gr.Blocks()
with gui:
gr.Markdown("# Live Video AI Assistant")
with gr.Row():
video_component = gr.Video()
threading.Thread(target=video_feed, daemon=True).start()
prompt = gr.Textbox(label="Enter your question")
response = gr.Textbox(label="AI Response")
btn = gr.Button("Ask")
btn.click(process_image, inputs=prompt, outputs=response)
gui.launch()