174 lines
7 KiB
Python
174 lines
7 KiB
Python
import os
|
|
import json
|
|
import logging
|
|
import torch
|
|
import cv2
|
|
import requests
|
|
import zipfile
|
|
import shutil
|
|
from ultralytics import YOLO
|
|
from urllib.parse import urlparse
|
|
|
|
def load_pipeline_node(node_config: dict, mpta_dir: str) -> dict:
|
|
# Recursively load a model node from configuration.
|
|
model_path = os.path.join(mpta_dir, node_config["modelFile"])
|
|
if not os.path.exists(model_path):
|
|
logging.error(f"Model file {model_path} not found.")
|
|
raise FileNotFoundError(f"Model file {model_path} not found.")
|
|
logging.info(f"Loading model for node {node_config['modelId']} from {model_path}")
|
|
model = YOLO(model_path)
|
|
if torch.cuda.is_available():
|
|
model.to("cuda")
|
|
node = {
|
|
"modelId": node_config["modelId"],
|
|
"modelFile": node_config["modelFile"],
|
|
"triggerClasses": node_config.get("triggerClasses", []),
|
|
"crop": node_config.get("crop", False),
|
|
"minConfidence": node_config.get("minConfidence", None),
|
|
"model": model,
|
|
"branches": []
|
|
}
|
|
for child in node_config.get("branches", []):
|
|
node["branches"].append(load_pipeline_node(child, mpta_dir))
|
|
return node
|
|
|
|
def load_pipeline_from_zip(zip_source: str, target_dir: str) -> dict:
|
|
os.makedirs(target_dir, exist_ok=True)
|
|
zip_path = os.path.join(target_dir, "pipeline.mpta")
|
|
|
|
# Parse the source; only local files are supported here.
|
|
parsed = urlparse(zip_source)
|
|
if parsed.scheme in ("", "file"):
|
|
local_path = parsed.path if parsed.scheme == "file" else zip_source
|
|
if os.path.exists(local_path):
|
|
try:
|
|
shutil.copy(local_path, zip_path)
|
|
logging.info(f"Copied local .mpta file from {local_path} to {zip_path}")
|
|
except Exception as e:
|
|
logging.error(f"Failed to copy local .mpta file from {local_path}: {e}")
|
|
return None
|
|
else:
|
|
logging.error(f"Local file {local_path} does not exist.")
|
|
return None
|
|
else:
|
|
logging.error("HTTP download functionality has been moved. Use a local file path here.")
|
|
return None
|
|
|
|
try:
|
|
with zipfile.ZipFile(zip_path, "r") as zip_ref:
|
|
zip_ref.extractall(target_dir)
|
|
logging.info(f"Extracted .mpta file to {target_dir}")
|
|
except Exception as e:
|
|
logging.error(f"Failed to extract .mpta file: {e}")
|
|
return None
|
|
finally:
|
|
if os.path.exists(zip_path):
|
|
os.remove(zip_path)
|
|
pipeline_name = os.path.basename(zip_source)
|
|
pipeline_name = os.path.splitext(pipeline_name)[0]
|
|
mpta_dir = os.path.join(target_dir, pipeline_name)
|
|
pipeline_json_path = os.path.join(mpta_dir, "pipeline.json")
|
|
if not os.path.exists(pipeline_json_path):
|
|
logging.error("pipeline.json not found in the .mpta file")
|
|
return None
|
|
|
|
try:
|
|
with open(pipeline_json_path, "r") as f:
|
|
pipeline_config = json.load(f)
|
|
return load_pipeline_node(pipeline_config["pipeline"], mpta_dir)
|
|
except Exception as e:
|
|
logging.error(f"Error loading pipeline.json: {e}")
|
|
return None
|
|
|
|
def run_pipeline(frame, node: dict, return_bbox: bool = False):
|
|
"""
|
|
Processes the frame with the given pipeline node. When return_bbox is True,
|
|
the function returns a tuple (detection, bbox) where bbox is (x1,y1,x2,y2)
|
|
for drawing. Otherwise, returns only the detection.
|
|
"""
|
|
try:
|
|
# Check model type and use appropriate method
|
|
model_task = getattr(node["model"], "task", None)
|
|
|
|
if model_task == "classify":
|
|
# Classification models need to use predict() instead of track()
|
|
logging.debug(f"Running classification model: {node.get('modelId')}")
|
|
results = node["model"].predict(frame, stream=False)
|
|
detection = None
|
|
best_box = None
|
|
|
|
# Process classification results
|
|
for r in results:
|
|
probs = r.probs
|
|
if probs is not None and len(probs) > 0:
|
|
# Get the most confident class
|
|
class_id = int(probs.top1)
|
|
conf = float(probs.top1conf)
|
|
detection = {
|
|
"class": node["model"].names[class_id],
|
|
"confidence": conf,
|
|
"id": None # Classification doesn't have tracking IDs
|
|
}
|
|
|
|
# Classification doesn't produce bounding boxes
|
|
bbox = None
|
|
|
|
else:
|
|
# Detection/segmentation models use tracking
|
|
logging.debug(f"Running detection/tracking model: {node.get('modelId')}")
|
|
results = node["model"].track(frame, stream=False, persist=True)
|
|
detection = None
|
|
best_box = None
|
|
max_conf = -1
|
|
|
|
for r in results:
|
|
for box in r.boxes:
|
|
box_cpu = box.cpu()
|
|
conf = float(box_cpu.conf[0])
|
|
if conf > max_conf and hasattr(box, "id") and box.id is not None:
|
|
max_conf = conf
|
|
detection = {
|
|
"class": node["model"].names[int(box_cpu.cls[0])],
|
|
"confidence": conf,
|
|
"id": box.id.item()
|
|
}
|
|
best_box = box_cpu
|
|
|
|
bbox = None
|
|
# Calculate bbox if best_box exists
|
|
if detection and best_box is not None:
|
|
coords = best_box.xyxy[0]
|
|
x1, y1, x2, y2 = map(int, coords)
|
|
h, w = frame.shape[:2]
|
|
x1, y1 = max(0, x1), max(0, y1)
|
|
x2, y2 = min(w, x2), min(h, y2)
|
|
if x2 > x1 and y2 > y1:
|
|
bbox = (x1, y1, x2, y2)
|
|
if node.get("crop", False):
|
|
frame = frame[y1:y2, x1:x2]
|
|
|
|
if detection is not None:
|
|
for branch in node["branches"]:
|
|
if detection["class"] in branch.get("triggerClasses", []):
|
|
min_conf = branch.get("minConfidence")
|
|
if min_conf is not None and detection["confidence"] < min_conf:
|
|
logging.debug(f"Confidence {detection['confidence']} below threshold {min_conf} for branch {branch['modelId']}.")
|
|
if return_bbox:
|
|
return detection, bbox
|
|
return detection
|
|
res = run_pipeline(frame, branch, return_bbox)
|
|
if res is not None:
|
|
if return_bbox:
|
|
return res
|
|
return res
|
|
if return_bbox:
|
|
return detection, bbox
|
|
return detection
|
|
if return_bbox:
|
|
return None, None
|
|
return None
|
|
except Exception as e:
|
|
logging.error(f"Error running pipeline on node {node.get('modelId')}: {e}")
|
|
if return_bbox:
|
|
return None, None
|
|
return None
|