""" Простой скрипт для детекции печатей (stamp) Требуется только: модель и изображение """ import cv2 import os import sys import json from ultralytics import YOLO def detect_stamps_no_save(image_path, model_path="stamp_model.pt", conf=0.25, model=None): """ Detect stamps without saving images. Args: image_path: Path to input image model_path: Path to model (or will download from HF Hub if not found) conf: Confidence threshold model: Pre-loaded YOLO model (optional, will load if not provided) Returns: dict: Detection results with detections and image_size """ # Use pre-loaded model if provided, otherwise load model if model is None: # Load model - try to download from HF Hub if not found locally if not os.path.exists(model_path): # Try to download from Hugging Face Hub try: from huggingface_hub import hf_hub_download print(f"Model not found locally, attempting to download from HF Hub...") # You can upload your model to HF Hub and use it here # For now, try the default path in stamp_detector directory default_path = os.path.join("stamp_detector", "stamp_model.pt") if os.path.exists(default_path): model_path = default_path else: raise FileNotFoundError(f"Stamp model not found: {model_path}. Please upload stamp_model.pt to the Space.") except ImportError: raise FileNotFoundError(f"Stamp model not found: {model_path}") model = YOLO(model_path) # Load image if not os.path.exists(image_path): raise FileNotFoundError(f"Image not found: {image_path}") image = cv2.imread(image_path) if image is None: raise ValueError(f"Could not load image: {image_path}") # Detection results = model(image, conf=conf, verbose=False) # Collect detections detections = [] image_height, image_width = image.shape[:2] for result in results: boxes = result.boxes for box in boxes: class_id = int(box.cls[0]) confidence = float(box.conf[0]) # Filter only stamp (class_id == 0) if class_id == 0 and confidence >= conf: x1, y1, x2, y2 = map(int, box.xyxy[0]) detection = { "class": "stamp", "confidence": round(confidence, 4), "bbox": { "x1": int(x1), "y1": int(y1), "x2": int(x2), "y2": int(y2), "width": int(x2 - x1), "height": int(y2 - y1) }, "bbox_normalized": { "x1": round(x1 / image_width, 6), "y1": round(y1 / image_height, 6), "x2": round(x2 / image_width, 6), "y2": round(y2 / image_height, 6), "width": round((x2 - x1) / image_width, 6), "height": round((y2 - y1) / image_height, 6) } } detections.append(detection) return { "image_size": { "width": image_width, "height": image_height }, "detections_count": len(detections), "detections": detections } def detect_stamps(image_path, model_path="stamp_model.pt", output_path=None, conf=0.25, return_json=False): """ Детектирует печати на изображении Args: image_path: путь к входному изображению model_path: путь к модели (по умолчанию: stamp_model.pt) output_path: путь для сохранения результата (если None, создается автоматически) conf: порог уверенности (по умолчанию: 0.25) return_json: если True, возвращает также JSON с координатами Returns: если return_json=False: путь к сохраненному изображению если return_json=True: словарь с 'image_path' и 'detections' (JSON структура) """ # Загружаем модель if not os.path.exists(model_path): print(f"❌ Ошибка: модель не найдена: {model_path}") sys.exit(1) print(f"📥 Загружаю модель: {model_path}") model = YOLO(model_path) print("✅ Модель загружена") # Загружаем изображение if not os.path.exists(image_path): print(f"❌ Ошибка: изображение не найдено: {image_path}") sys.exit(1) print(f"📷 Загружаю изображение: {image_path}") image = cv2.imread(image_path) if image is None: print(f"❌ Ошибка: не удалось загрузить изображение") sys.exit(1) # Детекция print(f"🔍 Выполняю детекцию (порог: {conf})...") results = model(image, conf=conf, verbose=False) # Собираем детекции и рисуем рамки result_image = image.copy() detections = [] image_height, image_width = image.shape[:2] for result in results: boxes = result.boxes for box in boxes: class_id = int(box.cls[0]) confidence = float(box.conf[0]) # Фильтруем только stamp (class_id == 0) if class_id == 0 and confidence >= conf: x1, y1, x2, y2 = map(int, box.xyxy[0]) # Сохраняем детекцию в JSON формате detection = { "class": "stamp", "confidence": round(confidence, 4), "bbox": { "x1": int(x1), "y1": int(y1), "x2": int(x2), "y2": int(y2), "width": int(x2 - x1), "height": int(y2 - y1) }, "bbox_normalized": { "x1": round(x1 / image_width, 6), "y1": round(y1 / image_height, 6), "x2": round(x2 / image_width, 6), "y2": round(y2 / image_height, 6), "width": round((x2 - x1) / image_width, 6), "height": round((y2 - y1) / image_height, 6) } } detections.append(detection) # Рисуем рамку (красная) cv2.rectangle(result_image, (x1, y1), (x2, y2), (0, 0, 255), 2) # Подпись label = f"stamp {confidence:.2f}" (label_width, label_height), _ = cv2.getTextSize( label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2 ) cv2.rectangle( result_image, (x1, y1 - label_height - 10), (x1 + label_width, y1), (0, 0, 255), -1 ) cv2.putText( result_image, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2 ) # Сохраняем результат if output_path is None: base_name = os.path.splitext(os.path.basename(image_path))[0] output_dir = "output" os.makedirs(output_dir, exist_ok=True) output_path = os.path.join(output_dir, f"{base_name}_result.jpg") cv2.imwrite(output_path, result_image) print(f"✅ Найдено печатей: {len(detections)}") print(f"📁 Результат сохранен: {output_path}") # Возвращаем результат if return_json: result_data = { "image_path": output_path, "image_size": { "width": image_width, "height": image_height }, "detections_count": len(detections), "detections": detections } return result_data else: return output_path if __name__ == "__main__": import argparse parser = argparse.ArgumentParser( description="Детекция печатей на изображениях") parser.add_argument("image_path", help="Путь к изображению") parser.add_argument( "--model", default="stamp_model.pt", help="Путь к модели (по умолчанию: stamp_model.pt)" ) parser.add_argument( "--output", default=None, help="Путь для сохранения результата (по умолчанию: output/{имя_файла}_result.jpg)" ) parser.add_argument( "--conf", type=float, default=0.25, help="Порог уверенности (по умолчанию: 0.25)" ) parser.add_argument( "--json", action="store_true", help="Сохранить JSON с координатами детекций" ) parser.add_argument( "--json-output", default=None, help="Путь для сохранения JSON файла (по умолчанию: output/{имя_файла}_result.json)" ) args = parser.parse_args() print("=" * 60) print("🔍 Детекция печатей (stamp)") print("=" * 60) result = detect_stamps( args.image_path, args.model, args.output, args.conf, return_json=args.json or args.json_output is not None ) # Сохраняем JSON если нужно if args.json or args.json_output is not None: if isinstance(result, dict): json_data = { "image_path": result["image_path"], "image_size": result["image_size"], "detections_count": result["detections_count"], "detections": result["detections"] } else: # Если result - это путь, нужно пересчитать result = detect_stamps( args.image_path, args.model, args.output, args.conf, return_json=True ) json_data = { "image_path": result["image_path"], "image_size": result["image_size"], "detections_count": result["detections_count"], "detections": result["detections"] } # Определяем путь для JSON if args.json_output: json_path = args.json_output else: base_name = os.path.splitext(os.path.basename(args.image_path))[0] output_dir = "output" os.makedirs(output_dir, exist_ok=True) json_path = os.path.join(output_dir, f"{base_name}_result.json") # Сохраняем JSON with open(json_path, "w", encoding="utf-8") as f: json.dump(json_data, f, indent=2, ensure_ascii=False) print(f"📄 JSON сохранен: {json_path}") print("=" * 60)