Update Lib/Core.py
Browse files- Lib/Core.py +8 -1
Lib/Core.py
CHANGED
|
@@ -8,11 +8,18 @@ import torch
|
|
| 8 |
from matplotlib import pyplot as plt
|
| 9 |
from ultralytics import YOLO
|
| 10 |
from Lib.Consts import LABELS, COLOR_MAP, COLOR_MAP_RGB
|
| 11 |
-
|
| 12 |
|
| 13 |
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
| 14 |
|
| 15 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 16 |
class CablePoleSegmentation():
|
| 17 |
def __init__(self, model_path=None, retina_mask=False):
|
| 18 |
if not model_path:
|
|
|
|
| 8 |
from matplotlib import pyplot as plt
|
| 9 |
from ultralytics import YOLO
|
| 10 |
from Lib.Consts import LABELS, COLOR_MAP, COLOR_MAP_RGB
|
| 11 |
+
from Tool.Core import DownloadHFModel
|
| 12 |
|
| 13 |
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
| 14 |
|
| 15 |
|
| 16 |
+
# Download Model
|
| 17 |
+
REPO_ID = "blitzkrieg0000/yolov9_pole-cable-detection"
|
| 18 |
+
MODEL_FILE = "yolov9c-cable-seg.pt"
|
| 19 |
+
DownloadHFModel(REPO_ID, MODEL_FILE)
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
|
| 23 |
class CablePoleSegmentation():
|
| 24 |
def __init__(self, model_path=None, retina_mask=False):
|
| 25 |
if not model_path:
|