Spaces:
Sleeping
Sleeping
File size: 3,693 Bytes
e56eb98 8f756c5 e56eb98 8f756c5 e56eb98 8f756c5 e56eb98 | 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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | from pathlib import Path
import cv2
import numpy as np
from ultralytics import YOLO
import os
os.environ['YOLO_VERBOSE'] = 'False'
# Module-level cache: avoids reloading YOLO weights on every detect_logo call.
# Key: model_path string → Value: YOLO instance
_YOLO_CACHE: dict = {}
def _get_yolo(model_path: str) -> YOLO:
"""Return a cached YOLO instance, loading from disk only on first call."""
if model_path not in _YOLO_CACHE:
_YOLO_CACHE[model_path] = YOLO(model_path, verbose=False)
return _YOLO_CACHE[model_path]
def ensure_portrait(image: np.ndarray) -> np.ndarray:
if image is None or image.size == 0:
return image
height, width = image.shape[:2]
if width > height:
return cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
return image
def enhance_contrast(image: np.ndarray, clip_limit: float) -> np.ndarray:
if image is None or image.size == 0:
return image
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(8, 8))
gray_eq = clahe.apply(gray)
return cv2.cvtColor(gray_eq, cv2.COLOR_GRAY2BGR)
def detect_logo(image_path: str, model_path: str = "models\logo.pt") -> int:
"""
Detect logo in image and return class ID.
Args:
image_path: Path to input image
model_path: Path to YOLO model weights
Returns:
0: Uttarakhand
1: CBSE
2: ICSE
-1: No detection
"""
model = _get_yolo(model_path)
img = cv2.imread(image_path)
if img is None:
return -1
img = ensure_portrait(img)
for clip_val in range(9, 14):
processed = enhance_contrast(img, float(clip_val))
results = model.predict(source=processed, verbose=False)
if results and results[0].boxes is not None and len(results[0].boxes) > 0:
conf = results[0].boxes.conf.cpu().numpy()
cls = results[0].boxes.cls.cpu().numpy()
valid_mask = conf >= 0.25
if valid_mask.any():
best_idx = int(np.argmax(conf[valid_mask]))
return int(cls[valid_mask][best_idx])
return -1
def detect_logo_with_boxes(image_path: str, model_path: str = "models\logo.pt") -> list:
"""
Detect logo in image and return bounding boxes.
Args:
image_path: Path to input image
model_path: Path to YOLO model weights
Returns:
list: List of bounding boxes [(x1, y1, x2, y2), ...]
"""
print(f"Logo detection with boxes for: {image_path}")
model = _get_yolo(model_path)
img = cv2.imread(image_path)
if img is None:
print(f"Failed to load image: {image_path}")
return []
print(f"Image loaded, shape: {img.shape}")
img = ensure_portrait(img)
print(f"After portrait check, shape: {img.shape}")
for clip_val in range(9, 14):
processed = enhance_contrast(img, float(clip_val))
results = model.predict(source=processed, verbose=False)
if results and results[0].boxes is not None and len(results[0].boxes) > 0:
conf = results[0].boxes.conf.cpu().numpy()
cls = results[0].boxes.cls.cpu().numpy()
boxes = results[0].boxes.xyxy.cpu().numpy()
print(f"Found {len(conf)} detections with confidences: {conf}")
valid_mask = conf >= 0.25
if valid_mask.any():
# Return the best detection (highest confidence) like the original function
best_idx = int(np.argmax(conf[valid_mask]))
best_box = boxes[valid_mask][best_idx]
print(f"Best detection box: {best_box}")
return [best_box.tolist()]
print("No valid logo detections found")
return []
if __name__ == "__main__":
import sys
if len(sys.argv) < 2:
print("Usage: python detectLogo.py <image_path> [model_path]")
sys.exit(1)
image_path = sys.argv[1]
model_path = sys.argv[2] if len(sys.argv) > 2 else "models\logo.pt"
result = detect_logo(image_path, model_path)
print(result) |