new logic
This commit is contained in:
Pongsatorn
parent
192b96d658
commit
aa4e0463d4
4 changed files with 303 additions and 129 deletions
|
|
@ -3,172 +3,180 @@ 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}")
|
||||
logging.info(f"Loading model {node_config['modelId']} from {model_path}")
|
||||
model = YOLO(model_path)
|
||||
if torch.cuda.is_available():
|
||||
model.to("cuda")
|
||||
node = {
|
||||
|
||||
# map triggerClasses names → indices for YOLO
|
||||
names = model.names # idx -> class name
|
||||
trigger_names = node_config.get("triggerClasses", [])
|
||||
trigger_inds = [i for i, nm in names.items() if nm in trigger_names]
|
||||
|
||||
return {
|
||||
"modelId": node_config["modelId"],
|
||||
"modelFile": node_config["modelFile"],
|
||||
"triggerClasses": node_config.get("triggerClasses", []),
|
||||
"triggerClasses": trigger_names,
|
||||
"triggerClassIndices": trigger_inds,
|
||||
"crop": node_config.get("crop", False),
|
||||
"minConfidence": node_config.get("minConfidence", None),
|
||||
"minConfidence": node_config.get("minConfidence", 0.0),
|
||||
"model": model,
|
||||
"branches": []
|
||||
"branches": [
|
||||
load_pipeline_node(child, mpta_dir)
|
||||
for child in node_config.get("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.")
|
||||
local = parsed.path if parsed.scheme == "file" else zip_source
|
||||
if not os.path.exists(local):
|
||||
logging.error(f"Local file {local} does not exist.")
|
||||
return None
|
||||
shutil.copy(local, zip_path)
|
||||
else:
|
||||
logging.error("HTTP download functionality has been moved. Use a local file path here.")
|
||||
logging.error("HTTP download not supported; use local file.")
|
||||
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")
|
||||
with zipfile.ZipFile(zip_path, "r") as z:
|
||||
z.extractall(target_dir)
|
||||
os.remove(zip_path)
|
||||
|
||||
base = os.path.splitext(os.path.basename(zip_source))[0]
|
||||
mpta_dir = os.path.join(target_dir, base)
|
||||
cfg = os.path.join(mpta_dir, "pipeline.json")
|
||||
if not os.path.exists(cfg):
|
||||
logging.error("pipeline.json not found in archive.")
|
||||
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
|
||||
with open(cfg) as f:
|
||||
pipeline_config = json.load(f)
|
||||
return load_pipeline_node(pipeline_config["pipeline"], mpta_dir)
|
||||
|
||||
def run_pipeline(frame, node: dict, return_bbox: bool = False):
|
||||
|
||||
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.
|
||||
- For detection nodes (task != 'classify'):
|
||||
• runs `track(..., classes=triggerClassIndices)`
|
||||
• picks top box ≥ minConfidence
|
||||
• optionally crops & resizes → recurse into child
|
||||
• else returns (det_dict, bbox)
|
||||
- For classify nodes:
|
||||
• runs `predict()`
|
||||
• returns top (class,confidence) and no bbox
|
||||
"""
|
||||
try:
|
||||
# Check model type and use appropriate method
|
||||
model_task = getattr(node["model"], "task", None)
|
||||
task = getattr(node["model"], "task", None)
|
||||
|
||||
# ─── Classification stage ───────────────────────────────────
|
||||
# if task == "classify":
|
||||
# results = node["model"].predict(frame, stream=False)
|
||||
# dets = []
|
||||
# for r in results:
|
||||
# probs = r.probs
|
||||
# if probs is not None:
|
||||
# # sort descending
|
||||
# idxs = probs.argsort(descending=True)
|
||||
# for cid in idxs:
|
||||
# dets.append({
|
||||
# "class": node["model"].names[int(cid)],
|
||||
# "confidence": float(probs[int(cid)]),
|
||||
# "id": None
|
||||
# })
|
||||
# if not dets:
|
||||
# return (None, None) if return_bbox else None
|
||||
|
||||
# best = dets[0]
|
||||
# return (best, None) if return_bbox else best
|
||||
|
||||
if model_task == "classify":
|
||||
# Classification models need to use predict() instead of track()
|
||||
logging.debug(f"Running classification model: {node.get('modelId')}")
|
||||
if task == "classify":
|
||||
# run the classifier and grab its top-1 directly via the Probs API
|
||||
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
|
||||
# nothing returned?
|
||||
if not results:
|
||||
return (None, None) if return_bbox else None
|
||||
|
||||
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
|
||||
# take the first result's probs object
|
||||
r = results[0]
|
||||
probs = r.probs
|
||||
if probs is None:
|
||||
return (None, None) if return_bbox else None
|
||||
|
||||
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]
|
||||
# get the top-1 class index and its confidence
|
||||
top1_idx = int(probs.top1)
|
||||
top1_conf = float(probs.top1conf)
|
||||
|
||||
det = {
|
||||
"class": node["model"].names[top1_idx],
|
||||
"confidence": top1_conf,
|
||||
"id": None
|
||||
}
|
||||
return (det, None) if return_bbox else det
|
||||
|
||||
|
||||
# ─── Detection stage ────────────────────────────────────────
|
||||
# only look for your triggerClasses
|
||||
tk = node["triggerClassIndices"]
|
||||
res = node["model"].track(
|
||||
frame,
|
||||
stream=False,
|
||||
persist=True,
|
||||
**({"classes": tk} if tk else {})
|
||||
)[0]
|
||||
|
||||
dets, boxes = [], []
|
||||
for box in res.boxes:
|
||||
conf = float(box.cpu().conf[0])
|
||||
cid = int(box.cpu().cls[0])
|
||||
name = node["model"].names[cid]
|
||||
if conf < node["minConfidence"]:
|
||||
continue
|
||||
xy = box.cpu().xyxy[0]
|
||||
x1,y1,x2,y2 = map(int, xy)
|
||||
dets.append({"class": name, "confidence": conf,
|
||||
"id": box.id.item() if hasattr(box, "id") else None})
|
||||
boxes.append((x1, y1, x2, y2))
|
||||
|
||||
if not dets:
|
||||
return (None, None) if return_bbox else None
|
||||
|
||||
# take highest‐confidence
|
||||
best_idx = max(range(len(dets)), key=lambda i: dets[i]["confidence"])
|
||||
best_det = dets[best_idx]
|
||||
best_box = boxes[best_idx]
|
||||
|
||||
# ─── Branch (classification) ───────────────────────────────
|
||||
for br in node["branches"]:
|
||||
if (best_det["class"] in br["triggerClasses"]
|
||||
and best_det["confidence"] >= br["minConfidence"]):
|
||||
# crop if requested
|
||||
sub = frame
|
||||
if br["crop"]:
|
||||
x1,y1,x2,y2 = best_box
|
||||
sub = frame[y1:y2, x1:x2]
|
||||
sub = cv2.resize(sub, (224, 224))
|
||||
|
||||
det2, _ = run_pipeline(sub, br, return_bbox=True)
|
||||
if det2:
|
||||
# return classification result + original bbox
|
||||
return (det2, best_box) if return_bbox else det2
|
||||
|
||||
# ─── No branch matched → return this detection ─────────────
|
||||
return (best_det, best_box) if return_bbox else best_det
|
||||
|
||||
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
|
||||
logging.error(f"Error in node {node.get('modelId')}: {e}")
|
||||
return (None, None) if return_bbox else None
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue