pympta & webcam dev tester
This commit is contained in:
parent
5da166a341
commit
ee0071284e
3 changed files with 190 additions and 155 deletions
158
app.py
158
app.py
|
|
@ -18,6 +18,9 @@ from fastapi.websockets import WebSocketDisconnect
|
|||
from websockets.exceptions import ConnectionClosedError
|
||||
from ultralytics import YOLO
|
||||
|
||||
# Import shared pipeline functions
|
||||
from siwatsystem.pympta import load_pipeline_from_zip, run_pipeline
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
# Global dictionaries to keep track of models and streams
|
||||
|
|
@ -57,161 +60,6 @@ WORKER_TIMEOUT_MS = 10000
|
|||
streams_lock = threading.Lock()
|
||||
models_lock = threading.Lock()
|
||||
|
||||
####################################################
|
||||
# Pipeline (Model)-loading helper functions
|
||||
####################################################
|
||||
def load_pipeline_node(node_config: dict, models_dir: str) -> dict:
|
||||
"""
|
||||
Recursively load a model node.
|
||||
Expects node_config to have:
|
||||
- modelId: a unique identifier
|
||||
- modelFile: the .pt file in models_dir
|
||||
- triggerClasses: list of class names that activate child branches
|
||||
- crop: boolean; if True, we crop to the bounding box for the next model
|
||||
- minConfidence: (optional) minimum confidence required to enter this branch
|
||||
- branches: list of child node configurations
|
||||
"""
|
||||
model_path = os.path.join(models_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), # NEW FIELD
|
||||
"model": model,
|
||||
"branches": []
|
||||
}
|
||||
for child_config in node_config.get("branches", []):
|
||||
child_node = load_pipeline_node(child_config, models_dir)
|
||||
node["branches"].append(child_node)
|
||||
return node
|
||||
|
||||
def load_pipeline_from_zip(zip_url: str, target_dir: str) -> dict:
|
||||
"""
|
||||
Download the .mpta file from zip_url, extract it to target_dir,
|
||||
and load the pipeline configuration (pipeline.json).
|
||||
Returns the model tree (root node) loaded with YOLO models.
|
||||
"""
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
zip_path = os.path.join(target_dir, "pipeline.mpta")
|
||||
|
||||
try:
|
||||
response = requests.get(zip_url, stream=True)
|
||||
if response.status_code == 200:
|
||||
with open(zip_path, "wb") as f:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
logging.info(f"Downloaded .mpta file from {zip_url} to {zip_path}")
|
||||
else:
|
||||
logging.error(f"Failed to download .mpta file (status {response.status_code})")
|
||||
return None
|
||||
except Exception as e:
|
||||
logging.error(f"Exception downloading .mpta file from {zip_url}: {e}")
|
||||
return None
|
||||
|
||||
# Extract the .mpta file
|
||||
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)
|
||||
|
||||
# Load pipeline.json
|
||||
pipeline_json_path = os.path.join(target_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)
|
||||
# Build the model tree recursively
|
||||
model_tree = load_pipeline_node(pipeline_config["pipeline"], target_dir)
|
||||
return model_tree
|
||||
except Exception as e:
|
||||
logging.error(f"Error loading pipeline.json: {e}")
|
||||
return None
|
||||
|
||||
####################################################
|
||||
# Model execution function
|
||||
####################################################
|
||||
def run_pipeline(frame, node: dict):
|
||||
"""
|
||||
Run the model at the current node.
|
||||
- Select the highest-confidence detection (if any).
|
||||
- If 'crop' is True, crop to the bounding box for the next stage.
|
||||
- If the detected class matches a branch's triggerClasses, check the confidence.
|
||||
If the detection's confidence is below branch["minConfidence"] (if specified),
|
||||
do not enter the branch and return the current detection.
|
||||
Returns the final detection result (dict) or None.
|
||||
"""
|
||||
try:
|
||||
results = node["model"].track(frame, stream=False, persist=True)
|
||||
detection = None
|
||||
max_conf = -1
|
||||
best_box = 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
|
||||
|
||||
# If there's a detection and crop is True, crop frame to bounding box
|
||||
if detection and node.get("crop", False) and best_box is not None:
|
||||
coords = best_box.xyxy[0] # [x1, y1, x2, y2]
|
||||
x1, y1, x2, y2 = map(int, coords)
|
||||
h, w = frame.shape[:2]
|
||||
x1 = max(0, x1)
|
||||
y1 = max(0, y1)
|
||||
x2 = min(w, x2)
|
||||
y2 = min(h, y2)
|
||||
|
||||
if x2 > x1 and y2 > y1:
|
||||
frame = frame[y1:y2, x1:x2] # crop the frame
|
||||
|
||||
if detection is not None:
|
||||
# Check if any branch should be entered based on trigger classes
|
||||
for branch in node["branches"]:
|
||||
if detection["class"] in branch.get("triggerClasses", []):
|
||||
# Check for a minimum confidence threshold for this branch
|
||||
min_conf = branch.get("minConfidence")
|
||||
if min_conf is not None and detection["confidence"] < min_conf:
|
||||
logging.debug(
|
||||
f"Detection confidence {detection['confidence']} below threshold "
|
||||
f"{min_conf} for branch {branch['modelId']}. Ending pipeline at current node."
|
||||
)
|
||||
return detection
|
||||
branch_detection = run_pipeline(frame, branch)
|
||||
if branch_detection is not None:
|
||||
return branch_detection
|
||||
return detection
|
||||
return None
|
||||
except Exception as e:
|
||||
logging.error(f"Error running pipeline on node {node.get('modelId')}: {e}")
|
||||
return None
|
||||
|
||||
####################################################
|
||||
# Detection and frame processing functions
|
||||
####################################################
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue