yasserrmd commited on
Commit
9c69830
1 Parent(s): 53d5146

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -6,7 +6,8 @@ import os
6
  import spaces
7
 
8
 
9
- model = YOLOv10.from_pretrained('jameslahm/yolov10x')
 
10
 
11
  # Define activity categories based on detected objects
12
  activity_categories = {
@@ -32,7 +33,6 @@ def categorize_activity(detected_objects):
32
  # Function to process the video, detect objects, and generate a categorized journal with images
33
  @spaces.GPU
34
  def generate_journal_with_images(video_path):
35
- model.to("cuda")
36
  cap = cv2.VideoCapture(video_path)
37
  journal_entries = {}
38
  saved_images = []
 
6
  import spaces
7
 
8
 
9
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
10
+ model = YOLOv10.from_pretrained('jameslahm/yolov10x').to(device)
11
 
12
  # Define activity categories based on detected objects
13
  activity_categories = {
 
33
  # Function to process the video, detect objects, and generate a categorized journal with images
34
  @spaces.GPU
35
  def generate_journal_with_images(video_path):
 
36
  cap = cv2.VideoCapture(video_path)
37
  journal_entries = {}
38
  saved_images = []