Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,073 Bytes
ce2b58f 5567efd 9538100 ce2b58f b0c7a24 bbc95d9 9538100 b0c7a24 8386bf1 b0c7a24 ce2b58f bbc95d9 ce2b58f d082050 ce2b58f bbc95d9 cdbafa3 ce2b58f bbc95d9 4cfdbcf b0c7a24 4cfdbcf b0c7a24 4cfdbcf 0bd515d b0c7a24 cdbafa3 bbc95d9 b0c7a24 e999761 b0c7a24 bbc95d9 b0c7a24 bbc95d9 b0c7a24 ce2b58f d082050 bbc95d9 d082050 bbc95d9 d082050 455843f ce2b58f b0c7a24 bbc95d9 ce2b58f bbc95d9 ce2b58f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import logging
from PIL import Image, ImageDraw
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
import os
import torch
# Setup logger
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
class ObjectDetector:
def __init__(self, model_key="yolov8n.pt", device="cpu"):
"""
Initializes an Ultralytics YOLO model path, defers actual model loading.
Args:
model_key (str): e.g. 'yolov8n.pt', 'yolov8s.pt', etc.
device (str): 'cpu' or 'cuda'
"""
self.device = device
resolved_key = model_key.lower().replace(".pt", "")
alias_map = {
"yolov8n": "yolov8n",
"yolov8s": "yolov8s",
"yolov8l": "yolov8l",
"yolov11b": "yolov11b"
}
hf_map = {
"yolov8n": ("ultralytics/yolov8", "yolov8n.pt"),
"yolov8s": ("ultralytics/yolov8", "yolov8s.pt"),
"yolov8l": ("ultralytics/yolov8", "yolov8l.pt"),
"yolov11b": ("Ultralytics/YOLO11", "yolov11b.pt"),
}
resolved_key = alias_map.get(resolved_key, resolved_key)
if resolved_key not in hf_map:
raise ValueError(f"Unsupported model key: {resolved_key}")
repo_id, filename = hf_map[resolved_key]
self.weights_path = hf_hub_download(
repo_id=repo_id,
filename=filename,
cache_dir="models/detection/weights",
force_download=False
)
logger.info(f"✅ YOLO weights ready for {resolved_key} at {self.weights_path}")
self.model = None # defer loading
def load_model(self):
if self.model is None:
logger.info("⚙️ Loading YOLO model into memory (runtime-safe)")
self.model = YOLO(self.weights_path)
if self.device == "cuda" and torch.cuda.is_available():
self.model.to("cuda")
logger.info(f"✅ YOLO model initialized on {self.device}")
return self
def predict(self, image: Image.Image, conf_threshold=0.25):
self.load_model()
logger.info("🔍 Running object detection")
results = self.model(image)
detections = []
for r in results:
for box in r.boxes:
detections.append({
"class_name": r.names[int(box.cls)],
"confidence": float(box.conf),
"bbox": box.xyxy[0].tolist()
})
logger.info(f"✅ Detected {len(detections)} objects")
return detections
def draw(self, image: Image.Image, detections, alpha=0.5):
overlay = image.copy()
draw = ImageDraw.Draw(overlay)
for det in detections:
bbox = det["bbox"]
label = f'{det["class_name"]} {det["confidence"]:.2f}'
draw.rectangle(bbox, outline="red", width=2)
draw.text((bbox[0], bbox[1]), label, fill="red")
return Image.blend(image, overlay, alpha)
|