Update sam_segment.py
Browse files- sam_segment.py +4 -3
sam_segment.py
CHANGED
|
@@ -6,8 +6,8 @@ import gc
|
|
| 6 |
|
| 7 |
# 定义可用的模型
|
| 8 |
MODELS = {
|
| 9 |
-
"small": "
|
| 10 |
-
"large": "
|
| 11 |
}
|
| 12 |
|
| 13 |
def clear_gpu_memory():
|
|
@@ -89,6 +89,7 @@ def segment_image_with_prompt(
|
|
| 89 |
|
| 90 |
# 准备模型参数
|
| 91 |
model_args = {
|
|
|
|
| 92 |
"retina_masks": True,
|
| 93 |
"conf": conf,
|
| 94 |
"iou": iou
|
|
@@ -124,7 +125,7 @@ def segment_image_with_prompt(
|
|
| 124 |
if hasattr(result, 'masks'):
|
| 125 |
del result.masks
|
| 126 |
del result
|
| 127 |
-
clear_gpu_memory()
|
| 128 |
|
| 129 |
return {
|
| 130 |
"total_segments": len(segments),
|
|
|
|
| 6 |
|
| 7 |
# 定义可用的模型
|
| 8 |
MODELS = {
|
| 9 |
+
"small": "./models/FastSAM-s.pt",
|
| 10 |
+
"large": "./models/FastSAM-x.pt"
|
| 11 |
}
|
| 12 |
|
| 13 |
def clear_gpu_memory():
|
|
|
|
| 89 |
|
| 90 |
# 准备模型参数
|
| 91 |
model_args = {
|
| 92 |
+
"device": "cpu",
|
| 93 |
"retina_masks": True,
|
| 94 |
"conf": conf,
|
| 95 |
"iou": iou
|
|
|
|
| 125 |
if hasattr(result, 'masks'):
|
| 126 |
del result.masks
|
| 127 |
del result
|
| 128 |
+
# clear_gpu_memory()
|
| 129 |
|
| 130 |
return {
|
| 131 |
"total_segments": len(segments),
|