Update app.py

app.py

@@ -1,38 +1,69 @@
-import streamlit as st
-from transformers import pipeline
-from io import BytesIO
-
-# Load Hugging Face models for question-answering and text-to-speech
-@st.cache_resource
-def load_qa_pipeline():
-    return pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")
-
-@st.cache_resource
-def load_tts_pipeline():
-    return pipeline("text-to-speech", model="facebook/fastspeech2-en-ljspeech")
-
-qa_pipeline = load_qa_pipeline()
-tts_pipeline = load_tts_pipeline()
-
-# Streamlit interface
-st.title("Virtual Assistant")
-st.write("Ask me anything!")
-
-# User query
-user_query = st.text_input("Type your question here:")
-
-if user_query:
-    # Generate answer using the QA model
-    context = "This is the context of the assistant. The assistant will answer general knowledge questions."  # Customize context for better QA accuracy
-    qa_result = qa_pipeline({"question": user_query, "context": context})
-    answer = qa_result['answer']
-    
-    # Display answer as text
-    st.write(f"Answer: {answer}")
-    
-    # Convert answer to audio using TTS model
-    tts_audio = tts_pipeline(answer, return_tensors="pt").audio
-    
-    # Streamlit audio player for TTS output
-    audio_bytes = BytesIO(tts_audio)
-    st.audio(audio_bytes, format="audio/wav")
+import gradio as gr
+import cv2
+import torch
+from transformers import DetrImageProcessor, DetrForObjectDetection
+from PIL import Image
+import tempfile
+import os
+
+# Load the DETR processor and model
+processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
+model = DetrForObjectDetection.from_pretrained("ArrayDice/Vehicle_Detection_Model_Zoom")
+
+# Function to process video and detect vehicles
+def detect_vehicles_in_video(video_path):
+    # Load video
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    # Temporary output video file
+    temp_video = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+    out = cv2.VideoWriter(temp_video.name, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
+
+    # Process each frame
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Convert frame to PIL for model processing
+        pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+
+        # Run object detection
+        inputs = processor(images=pil_image, return_tensors="pt")
+        outputs = model(**inputs)
+
+        # Set confidence threshold and draw boxes
+        for score, label, box in zip(outputs.logits.softmax(-1)[0], outputs.labels[0], outputs.boxes[0]):
+            if score.max() >= 0.9:  # Adjust confidence threshold as needed
+                x_min, y_min, x_max, y_max = map(int, box)
+                cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
+                label_text = f"Vehicle ({score.max():.2f})"
+                cv2.putText(frame, label_text, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+
+        # Write annotated frame to output video
+        out.write(frame)
+
+    # Release resources
+    cap.release()
+    out.release()
+
+    return temp_video.name  # Return path to the annotated video
+
+# Gradio interface
+def vehicle_detection_gradio(video):
+    annotated_video_path = detect_vehicles_in_video(video.name)
+    return annotated_video_path
+
+# Set up Gradio interface
+iface = gr.Interface(
+    fn=vehicle_detection_gradio,
+    inputs=gr.Video(type="file", label="Upload a video for vehicle detection"),
+    outputs=gr.Video(label="Annotated Video with Detected Vehicles"),
+    title="Vehicle Detection Model",
+    description="Upload a video, and the model will detect and annotate vehicles in each frame."
+)
+
+iface.launch()