feat: remove object mode pipeline (YOLO, SigLIP, TRAINING_MODE, OBJECT_CLASS)
winnow is a face recognition training tool. Object mode required manual file placement with no Frigate API, pulled in torch/torchvision/transformers/ ultralytics (~2 GB), and was architecturally misaligned with the project goal. Removed: - process_object_mode (YOLO inference), get_yolo_model - SigLIP model stack (get_siglip_model, get_object_embedding, batch variant) - entity_type branching throughout diversity, embeddings, jobs, executor - TRAINING_MODE and OBJECT_CLASS env vars - torch, torchvision, transformers, ultralytics dependencies - pytorch index entries from pyproject.toml - Object mode from README (Modes section, env var table, How It Works)
This commit is contained in:
@@ -1,4 +1,4 @@
|
||||
"""Image processing functions for cropping faces and objects."""
|
||||
"""Image processing functions for cropping faces."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
@@ -10,9 +10,6 @@ from .config import Config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Lazy singleton
|
||||
_yolo_model = None
|
||||
|
||||
|
||||
def _save_jpeg(img: Image.Image, path: str) -> None:
|
||||
if img.mode != "RGB":
|
||||
@@ -20,17 +17,6 @@ def _save_jpeg(img: Image.Image, path: str) -> None:
|
||||
img.save(path, format="JPEG")
|
||||
|
||||
|
||||
def get_yolo_model():
|
||||
"""Singleton for YOLO model."""
|
||||
global _yolo_model
|
||||
if _yolo_model is None:
|
||||
from ultralytics import YOLO
|
||||
|
||||
logger.info("Loading YOLOv9c model...")
|
||||
_yolo_model = YOLO("yolov9c.pt")
|
||||
return _yolo_model
|
||||
|
||||
|
||||
def align_face(img: Image.Image, landmarks: list[list[float]] | np.ndarray) -> Image.Image | None:
|
||||
"""Align face using 5-point landmarks to standard ArcFace input format (112x112).
|
||||
|
||||
@@ -170,49 +156,4 @@ def process_face_mode(
|
||||
return face_crop.size
|
||||
|
||||
|
||||
def process_object_mode(
|
||||
img: Image.Image,
|
||||
config: dict,
|
||||
output_dir: str,
|
||||
count: int,
|
||||
) -> bool:
|
||||
"""Detect and crop objects using YOLO."""
|
||||
try:
|
||||
model = get_yolo_model()
|
||||
target_class = config.get("object_class", "dog")
|
||||
import torch
|
||||
|
||||
if os.getenv("FORCE_CPU", "").lower() in ("true", "1", "yes"):
|
||||
device = "cpu"
|
||||
elif hasattr(torch, "xpu") and torch.xpu.is_available():
|
||||
device = "xpu"
|
||||
else:
|
||||
device = None # YOLO auto-selects (CUDA/ROCm/CPU)
|
||||
|
||||
results = model(img, verbose=False, device=device)
|
||||
|
||||
found = False
|
||||
class_idx = 0 # Sequential counter per target class (Issue #10)
|
||||
for box in (box for r in results for box in r.boxes):
|
||||
cls_id = int(box.cls[0])
|
||||
conf = float(box.conf[0])
|
||||
if 0 <= cls_id < len(model.names) and model.names[cls_id] == target_class and conf > 0.5:
|
||||
x1, y1, x2, y2 = box.xyxy[0].tolist()
|
||||
_save_jpeg(
|
||||
img.crop((x1, y1, x2, y2)),
|
||||
os.path.join(output_dir, f"{count}_{class_idx}.jpg"),
|
||||
)
|
||||
class_idx += 1
|
||||
found = True
|
||||
|
||||
return found
|
||||
except Exception as e:
|
||||
logger.error(f"YOLO processing failed: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def process_full_mode(img: Image.Image, output_dir: str, count: int) -> bool:
|
||||
"""Save full image."""
|
||||
_save_jpeg(img, os.path.join(output_dir, f"{count}.jpg"))
|
||||
return True
|
||||
|
||||
|
||||
Reference in New Issue
Block a user