Spaces:
Sleeping
Sleeping
Update app.py
Browse files
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 |
-
|
16 |
-
|
|
|
|
|
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()
|