nagasurendra commited on
Commit
045842d
·
verified ·
1 Parent(s): 28bbe05

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -3
app.py CHANGED
@@ -6,14 +6,17 @@ import numpy as np
6
  from ultralytics import YOLO
7
  import tempfile
8
  import torch
 
9
 
10
  # Set Ultralytics config path
11
  os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics'
12
 
13
  # Custom function to load model with trusted weights
14
  def load_model_with_trusted_weights(model_path):
15
- import ultralytics.nn.tasks
16
- with torch.serialization.safe_globals([ultralytics.nn.tasks.DetectionModel]):
 
 
17
  return YOLO(model_path)
18
 
19
  # Load both YOLO models
@@ -107,4 +110,4 @@ with gr.Blocks() as app:
107
  )
108
 
109
  if __name__ == "__main__":
110
- app.launch()
 
6
  from ultralytics import YOLO
7
  import tempfile
8
  import torch
9
+ import ultralytics.nn.tasks
10
 
11
  # Set Ultralytics config path
12
  os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics'
13
 
14
  # Custom function to load model with trusted weights
15
  def load_model_with_trusted_weights(model_path):
16
+ with torch.serialization.safe_globals([
17
+ ultralytics.nn.tasks.DetectionModel,
18
+ torch.nn.modules.container.Sequential
19
+ ]):
20
  return YOLO(model_path)
21
 
22
  # Load both YOLO models
 
110
  )
111
 
112
  if __name__ == "__main__":
113
+ app.launch()