assisTen / app.py
dioarafl's picture
Update app.py
a30ad6a verified
raw
history blame
4.02 kB
import cv2
import gradio as gr
import tempfile
import torch
from torchvision.models.detection import fasterrcnn_resnet50_fpn
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
import soundfile as sf
class FasterRCNNDetector:
def __init__(self):
self.model = fasterrcnn_resnet50_fpn(pretrained=True)
self.model.eval()
self.classes = [
"__background__", "person", "bicycle", "car", "motorcycle", "airplane", "bus",
"train", "truck", "boat", "traffic light", "fire hydrant", "N/A", "stop sign",
"parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
"elephant", "bear", "zebra", "giraffe", "N/A", "backpack", "umbrella", "N/A", "N/A",
"handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
"kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
"bottle", "N/A", "wine glass", "cup", "fork", "knife", "spoon", "bowl",
"banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza",
"donut", "cake", "chair", "couch", "potted plant", "bed", "N/A", "dining table",
"N/A", "N/A", "toilet", "N/A", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "N/A", "book",
"clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"
]
def detect_objects(self, image):
image_pil = Image.fromarray(image)
transform = transforms.Compose([transforms.ToTensor()])
image_tensor = transform(image_pil).unsqueeze(0)
with torch.no_grad():
prediction = self.model(image_tensor)
boxes = prediction[0]['boxes']
labels = prediction[0]['labels']
scores = prediction[0]['scores']
for box, label, score in zip(boxes, labels, scores):
box = [int(i) for i in box]
cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
cv2.putText(image, self.classes[label], (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 2)
return image
class JarvisModels:
def __init__(self):
self.client1 = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
self.model = deepspeech.Model("deepspeech-0.9.3-models.pbmm")
self.model.setBeamWidth(500)
async def generate_response(self, prompt):
generate_kwargs = dict(
temperature=0.6,
max_new_tokens=256,
top_p=0.95,
repetition_penalty=1,
do_sample=True,
seed=42,
)
formatted_prompt = system_instructions1 + prompt + "[JARVIS]"
stream = self.client1.text_generation(
formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=True)
output = ""
for response in stream:
output += response.token.text
communicate = edge_tts.Communicate(output)
with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
tmp_path = tmp_file.name
communicate.save(tmp_path)
return tmp_path
def transcribe_audio(audio_file):
model = JarvisModels().model
audio, sample_rate = sf.read(audio_file)
return model.stt(audio)
def generate_response(frame):
jarvis = JarvisModels()
response_model = await jarvis.generate_response("Hello, I see some interesting objects!")
return response_model
detector = FasterRCNNDetector()
iface = gr.Interface(
fn=[detector.detect_objects, transcribe_audio],
inputs=gr.inputs.Video(label="Webcam", parameters={"fps": 30}),
outputs=[gr.outputs.Image(), "text"],
title="Vision and Speech Interface",
description="This interface detects objects in the webcam feed and transcribes speech recorded through the microphone."
)
iface.launch()