Update app.py

app.py

@@ -1,47 +1,65 @@
 import gradio as gr
 from ultralytics import YOLOv10
 import cv2
+import torch
+import os
 import spaces
 
-# Load YOLOv10 model
-model = YOLOv10.from_pretrained('jameslahm/yolov10x')
+# Load YOLOv10 model with CUDA if available
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model = YOLOv10.from_pretrained('jameslahm/yolov10x').to(device)
 
-# Define object categories to classify activities
+# Define activity categories based on detected objects
 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
+    # Add more categories and objects as needed
 }
 
 # Function to map detected objects to categorized activities
 def categorize_activity(detected_objects):
-    activity_summary = {}
+    categorized_activities = {}
     
     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
+            if activity not in categorized_activities:
+                categorized_activities[activity] = []
+            categorized_activities[activity].append(detected_objects)
     
-    return activity_summary
+    return categorized_activities
 
-# Function to process the video and generate the journal
+# Function to process the video, detect objects, and generate a categorized journal with images
 @spaces.GPU
-def generate_journal(video):
-    cap = cv2.VideoCapture(video)
-    frame_rate = cap.get(cv2.CAP_PROP_FPS)
+def generate_journal_with_images(video_path):
+    cap = cv2.VideoCapture(video_path)
     journal_entries = {}
+    saved_images = []
+    frame_count = 0
+    output_folder = "detected_frames"
+    os.makedirs(output_folder, exist_ok=True)  # Create folder to store images
     
     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]
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        
+        # Make predictions using YOLOv10 on the current frame
+        results = model.predict(source=frame_rgb, device=device)
+        
+        # Draw bounding boxes on the frame
+        results.render()  # Render the results on the image (this modifies the frame in-place)
+        
+        # Save the image with bounding boxes
+        frame_filename = os.path.join(output_folder, f"frame_{frame_count}.jpg")
+        cv2.imwrite(frame_filename, frame_rgb[:, :, ::-1])  # Convert back to BGR for saving
+        saved_images.append(frame_filename)
+        
+        # Extract labels (class indices) and map them to class names
+        detected_objects = [model.names[int(box.cls)] for box in results.boxes]
         
         # Get current timestamp in the video
         timestamp = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000  # Convert ms to seconds
@@ -50,27 +68,41 @@ def generate_journal(video):
         activity_summary = categorize_activity(detected_objects)
         
         # Store the activities with their timestamp
-        for activity, count in activity_summary.items():
+        for activity, objects 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}")
+            journal_entries[activity].append((f"At {timestamp:.2f} seconds: {', '.join(objects[0])}", frame_filename))
+        
+        frame_count += 1
     
     cap.release()
     
-    # Create a formatted journal
+    # Create a formatted journal output
     formatted_journal = []
     for activity, entries in journal_entries.items():
         formatted_journal.append(f"**{activity}:**")
-        formatted_journal.extend(entries)
+        for entry, image_path in entries:
+            formatted_journal.append((entry, image_path))
     
-    return "\n".join(formatted_journal)
+    return 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"
-)
+# Gradio interface for uploading video and generating journal with images
+def display_journal_with_images(video):
+    journal_with_images = generate_journal_with_images(video)
+    
+    # Create the final display with text and images
+    display_items = []
+    for entry, image_path in journal_with_images:
+        display_items.append((entry, image_path))
+    
+    return display_items
+
+# Define Gradio Blocks for custom display
+with gr.Blocks() as iface:
+    video_input = gr.Video(label="Upload Video")
+    output_gallery = gr.Gallery(label="Generated Daily Journal with Images").style(grid=[2], height='auto')
+    run_button = gr.Button("Generate Journal")
+    
+    run_button.click(fn=display_journal_with_images, inputs=video_input, outputs=output_gallery)
 
 iface.launch()