DailySnap / app.py
yasserrmd's picture
Create app.py
100b9b2 verified
raw
history blame
2.59 kB
import gradio as gr
from ultralytics import YOLOv10
import cv2
import spaces
# Load YOLOv10 model
model = YOLOv10.from_pretrained('jameslahm/yolov10x')
# Define object categories to classify activities
activity_categories = {
"Working": ["laptop", "computer", "keyboard", "office chair"],
"Meal Time": ["fork", "spoon", "plate", "food"],
"Exercise": ["dumbbell", "bicycle", "yoga mat", "treadmill"],
"Outdoors": ["car", "tree", "bicycle", "road"],
# Add more categories and associated objects as needed
}
# Function to map detected objects to categorized activities
def categorize_activity(detected_objects):
activity_summary = {}
for activity, objects in activity_categories.items():
if any(obj in detected_objects for obj in objects):
if activity not in activity_summary:
activity_summary[activity] = 0
activity_summary[activity] += 1 # Increase count for that activity
return activity_summary
# Function to process the video and generate the journal
@spaces.GPU
def generate_journal(video):
cap = cv2.VideoCapture(video)
frame_rate = cap.get(cv2.CAP_PROP_FPS)
journal_entries = {}
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# Make predictions using YOLOv10
results = model.predict(source=frame)
detected_objects = [res.name for res in results]
# Get current timestamp in the video
timestamp = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000 # Convert ms to seconds
# Categorize the detected objects into activities
activity_summary = categorize_activity(detected_objects)
# Store the activities with their timestamp
for activity, count in activity_summary.items():
if activity not in journal_entries:
journal_entries[activity] = []
journal_entries[activity].append(f"At {timestamp:.2f} seconds: {count} objects related to {activity}")
cap.release()
# Create a formatted journal
formatted_journal = []
for activity, entries in journal_entries.items():
formatted_journal.append(f"**{activity}:**")
formatted_journal.extend(entries)
return "\n".join(formatted_journal)
# Gradio interface for uploading video and generating journal
iface = gr.Interface(
fn=generate_journal,
inputs=gr.Video(label="Upload Video"),
outputs=gr.Textbox(label="Generated Daily Journal"),
title="AI-Powered Daily Journal"
)
iface.launch()