Add project files
This commit is contained in:
Suman
858
pipeline/engine.py
Normal file
858
pipeline/engine.py
Normal file
@@ -0,0 +1,858 @@
|
||||
"""
|
||||
Orchestrates the whole match pipeline:
|
||||
|
||||
startup (once): bg-remove every template + precompute features for all
|
||||
four methods, so a request only ever has to process the
|
||||
single uploaded image.
|
||||
|
||||
per request: bg-remove the upload, then run SIFT / ORB / SuperGlue /
|
||||
LoFTR one after another. They used to run concurrently
|
||||
(real wall-clock ~max() instead of sum()), but on an
|
||||
8GB card that let SuperGlue and LoFTR hold overlapping
|
||||
GPU allocations at once, which was enough on its own to
|
||||
tip into CUDA OOM. Running them sequentially caps peak
|
||||
GPU memory to whichever single method needs the most,
|
||||
at the cost of some latency.
|
||||
"""
|
||||
|
||||
import gc
|
||||
import logging
|
||||
import os
|
||||
import resource
|
||||
import shutil
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
|
||||
import cv2
|
||||
import torch
|
||||
|
||||
import config
|
||||
from pipeline import (bg_removal, classical, color, color_grid, deep,
|
||||
flower_count, sam3_client, shape_match, texture_match,
|
||||
utils, vase_compare, yolo_world)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
METHODS = ("SIFT", "ORB", "SuperGlue", "LoFTR")
|
||||
|
||||
_templates_meta = {} # name -> {"original_filename", "nobg_path"}
|
||||
_bootstrapped = False
|
||||
|
||||
# Serializes each request's pipeline end-to-end. Combined with running the
|
||||
# four methods sequentially (see below), only one method for one request is
|
||||
# ever doing heavy CPU/GPU work at any instant across the whole process.
|
||||
_pipeline_lock = threading.Lock()
|
||||
|
||||
|
||||
def bootstrap():
|
||||
"""Precompute everything template-side. Safe to call more than once;
|
||||
only does real work the first time."""
|
||||
global _bootstrapped
|
||||
if _bootstrapped:
|
||||
return
|
||||
|
||||
logger.info("Preparing templates (bg removal + feature precompute for "
|
||||
"SIFT / ORB / SuperGlue / LoFTR)...")
|
||||
|
||||
for fname in utils.list_template_files():
|
||||
name = utils.template_name(fname)
|
||||
src_path = os.path.join(config.TEMPLATE_IMAGES_DIR, fname)
|
||||
|
||||
rgba = bg_removal.remove_background_file(src_path, config.TEMPLATES_NOBG_CACHE)
|
||||
bgr, _, mask = bg_removal.split_rgba(rgba)
|
||||
|
||||
kp, des = classical.extract_features("SIFT", bgr, mask)
|
||||
classical.set_template_features("SIFT", name, kp, des)
|
||||
|
||||
kp, des = classical.extract_features("ORB", bgr, mask)
|
||||
classical.set_template_features("ORB", name, kp, des)
|
||||
|
||||
deep.set_template_superpoint(name, bgr, mask)
|
||||
deep.set_template_loftr(name, bgr, mask)
|
||||
color.set_template_color_data(name, bgr, mask)
|
||||
shape_match.set_template_shape_data(name, mask)
|
||||
texture_match.set_template_texture_data(name, bgr, mask)
|
||||
|
||||
_templates_meta[name] = {
|
||||
"original_filename": fname,
|
||||
"nobg_path": os.path.join(config.TEMPLATES_NOBG_CACHE, name + ".png"),
|
||||
}
|
||||
logger.info(" template ready: %s", name)
|
||||
|
||||
_bootstrapped = True
|
||||
logger.info("%d templates ready (device: %s).", len(_templates_meta), deep.DEVICE)
|
||||
|
||||
|
||||
def templates_meta():
|
||||
return _templates_meta
|
||||
|
||||
|
||||
def _run_sift(bgr, mask):
|
||||
kp, des = classical.extract_features("SIFT", bgr, mask)
|
||||
return classical.match_against_templates("SIFT", kp, des)
|
||||
|
||||
|
||||
def _run_orb(bgr, mask):
|
||||
kp, des = classical.extract_features("ORB", bgr, mask)
|
||||
return classical.match_against_templates("ORB", kp, des)
|
||||
|
||||
|
||||
def _run_superglue(bgr, mask):
|
||||
return deep.superglue_match_against_templates(bgr, mask)
|
||||
|
||||
|
||||
def _run_loftr(bgr, mask):
|
||||
return deep.loftr_match_against_templates(bgr, mask)
|
||||
|
||||
|
||||
_METHOD_RUNNERS = {
|
||||
"SIFT": _run_sift,
|
||||
"ORB": _run_orb,
|
||||
"SuperGlue": _run_superglue,
|
||||
"LoFTR": _run_loftr,
|
||||
}
|
||||
|
||||
|
||||
def _timed(method, fn, *args):
|
||||
"""Runs one method's matcher, catching any exception so a single
|
||||
failing method (e.g. a CUDA OOM on an unusual image) is logged and
|
||||
reported back as a failed card instead of taking the whole request
|
||||
down with it."""
|
||||
start = time.perf_counter()
|
||||
try:
|
||||
result = fn(*args)
|
||||
error = None
|
||||
except Exception as e:
|
||||
logger.exception("Method %s failed", method)
|
||||
result = []
|
||||
error = str(e)
|
||||
elapsed = time.perf_counter() - start
|
||||
return result, elapsed, error
|
||||
|
||||
|
||||
def _overall_best(method_results):
|
||||
"""Borda-count-style aggregate across methods: each method's ranking of
|
||||
templates contributes points, so one noisy method can't dominate the
|
||||
call the way a raw-score sum could."""
|
||||
points = {}
|
||||
for method, payload in method_results.items():
|
||||
ranked = payload["results"]
|
||||
n = len(ranked)
|
||||
for idx, row in enumerate(ranked):
|
||||
points[row["template"]] = points.get(row["template"], 0) + (n - idx)
|
||||
|
||||
if not points:
|
||||
return None
|
||||
return max(points, key=points.get)
|
||||
|
||||
|
||||
def _weighted_scores(method_results):
|
||||
"""Literal weighted sum per template: weighted_score(t) = sum_m
|
||||
WEIGHT[m] * raw_score(m, t). Returns (winning_template_or_None, ranked
|
||||
list of {"template", "weighted_score", "breakdown"} sorted descending).
|
||||
A method that errored contributes 0 everywhere (its results list is
|
||||
empty), it doesn't skew the total."""
|
||||
totals = {}
|
||||
breakdown = {} # template -> {method: contribution}
|
||||
|
||||
for method, payload in method_results.items():
|
||||
weight = config.METHOD_WEIGHTS.get(method, 0)
|
||||
for row in payload["results"]:
|
||||
template = row["template"]
|
||||
contribution = weight * row["score"]
|
||||
totals[template] = totals.get(template, 0.0) + contribution
|
||||
breakdown.setdefault(template, {})[method] = round(contribution, 2)
|
||||
|
||||
if not totals:
|
||||
return None, []
|
||||
|
||||
ranked = sorted(
|
||||
(
|
||||
{
|
||||
"template": template,
|
||||
"weighted_score": round(total, 2),
|
||||
"breakdown": breakdown.get(template, {}),
|
||||
}
|
||||
for template, total in totals.items()
|
||||
),
|
||||
key=lambda r: r["weighted_score"],
|
||||
reverse=True,
|
||||
)
|
||||
return ranked[0]["template"], ranked
|
||||
|
||||
|
||||
def _color_as_method_result(color_analysis, elapsed, error):
|
||||
"""Reshapes color.compare_input_to_templates()'s output into the same
|
||||
{"template", "score", "confidence"} shape the other methods use, so it
|
||||
can be handed to _weighted_scores() generically -- that's the only
|
||||
place this participates; it's deliberately kept out of the sequential
|
||||
METHODS loop (color analysis isn't a GPU/CPU-heavy per-template model
|
||||
call, no need for that machinery) and out of _overall_best (Borda stays
|
||||
exactly the original four methods, unchanged)."""
|
||||
results = [
|
||||
{"template": t["template"], "score": t["match_pct"], "confidence": t["match_pct"]}
|
||||
for t in color_analysis.get("templates", [])
|
||||
]
|
||||
best = results[0] if results else None
|
||||
return {
|
||||
"results": results,
|
||||
"time_sec": elapsed,
|
||||
"error": error,
|
||||
"best": best,
|
||||
"is_confident": bool(best and best["score"] >= config.SCORE_THRESHOLD["Color"]),
|
||||
"best_image_file": None,
|
||||
}
|
||||
|
||||
|
||||
def _build_family_grid(request_dir, input_bgr, input_mask, weighted_best):
|
||||
"""Purely visual, tied to the weighted-best template only: clusters
|
||||
both the upload and the winning template into color-region tiles,
|
||||
renders each as its own side-by-side grid image, and matches up
|
||||
corresponding regions by color to report an area-match % per region."""
|
||||
input_families = color_grid.cluster_families(input_bgr, input_mask)
|
||||
|
||||
template_rgba = cv2.imread(
|
||||
_templates_meta[weighted_best]["nobg_path"], cv2.IMREAD_UNCHANGED
|
||||
)
|
||||
template_bgr, _, template_mask = bg_removal.split_rgba(template_rgba)
|
||||
template_families = color_grid.cluster_families(template_bgr, template_mask)
|
||||
|
||||
matches = color_grid.match_families(input_families, template_families)
|
||||
overall_pct = color_grid.overall_area_match(matches)
|
||||
|
||||
input_grid_file = "input_family_grid.png"
|
||||
template_grid_file = "template_family_grid.png"
|
||||
cv2.imwrite(os.path.join(request_dir, input_grid_file),
|
||||
color_grid.render_family_grid(input_bgr, input_families))
|
||||
cv2.imwrite(os.path.join(request_dir, template_grid_file),
|
||||
color_grid.render_family_grid(template_bgr, template_families))
|
||||
|
||||
return {
|
||||
"template": weighted_best,
|
||||
"input_grid_file": input_grid_file,
|
||||
"template_grid_file": template_grid_file,
|
||||
"overall_area_match_pct": overall_pct,
|
||||
"matches": matches,
|
||||
"error": None,
|
||||
}
|
||||
|
||||
|
||||
def _build_shape_visuals(request_dir, input_shape_data, weighted_best):
|
||||
"""Purely visual, tied to the weighted-best template: renders both
|
||||
silhouettes (normalized into the same canvas) side by side plus an
|
||||
overlay highlighting exactly where they agree/diverge."""
|
||||
template_shape_data = shape_match._template_shape_data.get(weighted_best)
|
||||
if template_shape_data is None:
|
||||
return None
|
||||
|
||||
input_file = "input_silhouette.png"
|
||||
template_file = "template_silhouette.png"
|
||||
overlay_file = "shape_overlay.png"
|
||||
cv2.imwrite(os.path.join(request_dir, input_file),
|
||||
shape_match.render_silhouette(input_shape_data["canonical_mask"],
|
||||
shape_match._INPUT_COLOR))
|
||||
cv2.imwrite(os.path.join(request_dir, template_file),
|
||||
shape_match.render_silhouette(template_shape_data["canonical_mask"],
|
||||
shape_match._TEMPLATE_COLOR))
|
||||
cv2.imwrite(os.path.join(request_dir, overlay_file),
|
||||
shape_match.render_overlay(input_shape_data["canonical_mask"],
|
||||
template_shape_data["canonical_mask"]))
|
||||
|
||||
return {
|
||||
"input_file": input_file,
|
||||
"template_file": template_file,
|
||||
"overlay_file": overlay_file,
|
||||
}
|
||||
|
||||
|
||||
def _build_texture_visuals(request_dir, input_texture_data, weighted_best):
|
||||
"""Purely visual, tied to the weighted-best template: renders both
|
||||
images' LBP (local micro-pattern) maps side by side."""
|
||||
template_texture_data = texture_match._template_texture_data.get(weighted_best)
|
||||
if template_texture_data is None:
|
||||
return None
|
||||
|
||||
input_file = "input_lbp.png"
|
||||
template_file = "template_lbp.png"
|
||||
cv2.imwrite(os.path.join(request_dir, input_file),
|
||||
texture_match.render_lbp_visual(input_texture_data["lbp_image"],
|
||||
input_texture_data["mask_bool"]))
|
||||
cv2.imwrite(os.path.join(request_dir, template_file),
|
||||
texture_match.render_lbp_visual(template_texture_data["lbp_image"],
|
||||
template_texture_data["mask_bool"]))
|
||||
|
||||
return {
|
||||
"input_file": input_file,
|
||||
"template_file": template_file,
|
||||
}
|
||||
|
||||
|
||||
def _cleanup_old_uploads():
|
||||
if not os.path.isdir(config.UPLOADS_DIR):
|
||||
return
|
||||
cutoff = time.time() - config.MAX_UPLOAD_AGE_SECONDS
|
||||
for entry in os.listdir(config.UPLOADS_DIR):
|
||||
path = os.path.join(config.UPLOADS_DIR, entry)
|
||||
try:
|
||||
if os.path.isdir(path) and os.path.getmtime(path) < cutoff:
|
||||
shutil.rmtree(path, ignore_errors=True)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _peak_rss_mb():
|
||||
# ru_maxrss is KB on Linux, bytes on macOS -- this app only targets Linux.
|
||||
return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
|
||||
|
||||
|
||||
def _release_resources():
|
||||
"""Best-effort cleanup between requests. gc.collect() drops any
|
||||
lingering numpy/cv2/torch objects promptly instead of waiting for the
|
||||
next allocation to trigger a cycle; empty_cache() hands unused *cached*
|
||||
CUDA blocks back to the driver so nvidia-smi/other processes see them
|
||||
freed (it does not, and cannot, free host RAM)."""
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def _build_flower_count_comparison(input_count, template_count, template):
|
||||
diff = input_count - template_count
|
||||
if diff == 0:
|
||||
message = (
|
||||
f"Your upload and {template} both show {input_count} flowers -- "
|
||||
f"flower count doesn't look like a factor in the match score here."
|
||||
)
|
||||
else:
|
||||
more_fewer = "more" if diff > 0 else "fewer"
|
||||
message = (
|
||||
f"Your upload has {input_count} flowers detected vs {template_count} in "
|
||||
f"{template} -- {abs(diff)} {more_fewer}. No method above scores flower "
|
||||
f"count directly, but a mismatch like this usually also lowers the "
|
||||
f"keypoint (SIFT/ORB/SuperGlue/LoFTR) and shape/texture scores, since "
|
||||
f"there's simply more or less bouquet for them to match against -- so "
|
||||
f"it's often part of why the overall match isn't higher."
|
||||
)
|
||||
return {
|
||||
"input_count": input_count,
|
||||
"template_count": template_count,
|
||||
"diff": diff,
|
||||
"message": message,
|
||||
"error": None,
|
||||
}
|
||||
|
||||
|
||||
def count_flowers(request_id: str, template: str = None) -> dict:
|
||||
"""On-demand flower counting for an already-processed upload, via two
|
||||
independent methods shown side by side, plus (if a matched template
|
||||
name is given) a vase-identity comparison:
|
||||
|
||||
- SAM3 (instance segmentation): finds individual flower-shaped
|
||||
regions by prompting SAM3 with the plain-English concept "flower".
|
||||
Unlike SAM1's old "segment everything, then guess" approach, this
|
||||
never proposes the vase or any ribbon/bow in the first place -- no
|
||||
exclude-box filtering needed. Runs as a subprocess into a separate
|
||||
Python 3.10 env (see pipeline/sam3_client.py) since transformers'
|
||||
Sam3Model needs Python>=3.10, incompatible with this app's own env.
|
||||
The same subprocess call also fetches SAM3's "vase" mask for the
|
||||
upload (and the template's, if a match is known) for the vase
|
||||
comparison below -- one model load covers every job.
|
||||
- YOLO-World (open-vocabulary detection): independent second opinion,
|
||||
unchanged from before -- its own "flower"/"vase"/"ribbon"/"bow"
|
||||
counts and box visualization, shown as its own coarser corroborating
|
||||
signal (one box per contiguous flower region, not per bloom).
|
||||
- Vase comparison (DINOv2 + CLIP): crops the vase out of the upload
|
||||
and the matched template using SAM3's precise masks (background
|
||||
pixels inside the crop are blacked out, so the embeddings only see
|
||||
the vase itself), then embeds both crops with DINOv2 (fine detail)
|
||||
and CLIP (holistic/semantic) and reports a same/uncertain/different
|
||||
verdict. Skipped (not an error) if no vase was detected on either
|
||||
side.
|
||||
|
||||
Heavyweight and strictly opt-in -- triggered only by its own UI button,
|
||||
never run as part of process_upload. Unloads the matching pipeline's own
|
||||
GPU-resident models first (SuperPoint/LightGlue/LoFTR, rembg's
|
||||
onnxruntime session) so YOLO-World/DINOv2/CLIP never have to share the
|
||||
8GB card with them (SAM3 runs in its own process/env regardless), then
|
||||
unloads YOLO-World/DINOv2/CLIP again afterward -- everything lazily
|
||||
reloads itself on whichever request needs it next, exactly like a fresh
|
||||
process start would."""
|
||||
request_dir = os.path.join(config.UPLOADS_DIR, request_id)
|
||||
nobg_path = os.path.join(request_dir, "nobg.png")
|
||||
if not os.path.isfile(nobg_path):
|
||||
raise FileNotFoundError(f"no processed upload found for request {request_id}")
|
||||
|
||||
with _pipeline_lock:
|
||||
start = time.perf_counter()
|
||||
rgba = cv2.imread(nobg_path, cv2.IMREAD_UNCHANGED)
|
||||
bgr, _, mask = bg_removal.split_rgba(rgba)
|
||||
|
||||
freed_deep = deep.unload_models()
|
||||
freed_rembg = bg_removal.unload_session()
|
||||
if freed_deep or freed_rembg:
|
||||
logger.info("[%s] unloaded matching-pipeline models before "
|
||||
"SAM3/YOLO-World/DINO/CLIP run", request_id)
|
||||
|
||||
try:
|
||||
by_class = yolo_world.detect(bgr)
|
||||
yolo_error = None
|
||||
except Exception as e:
|
||||
logger.exception("[%s] YOLO-World detection failed", request_id)
|
||||
by_class = {}
|
||||
yolo_error = str(e)
|
||||
|
||||
template_bgr = None
|
||||
template_mask = None
|
||||
if template is not None and template in _templates_meta:
|
||||
template_rgba = cv2.imread(
|
||||
_templates_meta[template]["nobg_path"], cv2.IMREAD_UNCHANGED
|
||||
)
|
||||
template_bgr, _, template_mask = bg_removal.split_rgba(template_rgba)
|
||||
|
||||
sam3_images = {"input": bgr}
|
||||
sam3_jobs = [
|
||||
{"image": "input", "prompt": config.SAM3_FLOWER_PROMPT,
|
||||
"threshold": config.SAM3_FLOWER_THRESHOLD},
|
||||
{"image": "input", "prompt": config.SAM3_VASE_PROMPT,
|
||||
"threshold": config.SAM3_VASE_THRESHOLD},
|
||||
]
|
||||
if template_bgr is not None:
|
||||
sam3_images["template"] = template_bgr
|
||||
sam3_jobs.append({"image": "template", "prompt": config.SAM3_VASE_PROMPT,
|
||||
"threshold": config.SAM3_VASE_THRESHOLD})
|
||||
# Also count flowers in the template photo itself (same "flower"
|
||||
# concept prompt) -- lets the UI tell the user when a lower match
|
||||
# score might partly be explained by a different flower count,
|
||||
# rather than leaving that as an unexplained low number.
|
||||
sam3_jobs.append({"image": "template", "prompt": config.SAM3_FLOWER_PROMPT,
|
||||
"threshold": config.SAM3_FLOWER_THRESHOLD})
|
||||
|
||||
try:
|
||||
sam3_result = sam3_client.run_jobs(sam3_images, sam3_jobs, workdir=request_dir)
|
||||
sam3_error = None
|
||||
except Exception as e:
|
||||
logger.exception("[%s] SAM3 failed", request_id)
|
||||
sam3_result = {}
|
||||
sam3_error = str(e)
|
||||
|
||||
try:
|
||||
flower_instances = sam3_result.get(("input", config.SAM3_FLOWER_PROMPT), [])
|
||||
count_data = flower_count.count_flowers(bgr, mask, request_dir,
|
||||
instances_raw=flower_instances)
|
||||
sam_error = sam3_error
|
||||
except Exception as e:
|
||||
logger.exception("[%s] flower counting (SAM3) failed", request_id)
|
||||
count_data = {"total_count": 0, "clusters": []}
|
||||
sam_error = sam3_error or str(e)
|
||||
|
||||
sam_visual_file = None
|
||||
if sam_error is None:
|
||||
try:
|
||||
visual = flower_count.render_instances(bgr, count_data)
|
||||
sam_visual_file = "flower_count_sam.png"
|
||||
cv2.imwrite(os.path.join(request_dir, sam_visual_file), visual)
|
||||
except Exception:
|
||||
logger.exception("[%s] SAM3 flower count visualization failed", request_id)
|
||||
|
||||
# Same flower-instance segmentation, run on the matched template's
|
||||
# own photo -- lets callers (the web UI's mismatch note, tester.py's
|
||||
# report image) show/compare "your photo's flowers" side by side
|
||||
# with "the template's flowers", not just two bare numbers.
|
||||
template_count_data = {"total_count": 0, "clusters": []}
|
||||
template_sam_visual_file = None
|
||||
template_sam_error = None
|
||||
if template_bgr is not None and template_mask is not None:
|
||||
if sam3_error is not None:
|
||||
template_sam_error = sam3_error
|
||||
else:
|
||||
try:
|
||||
template_flower_instances = sam3_result.get(
|
||||
("template", config.SAM3_FLOWER_PROMPT), []
|
||||
)
|
||||
template_count_data = flower_count.count_flowers(
|
||||
template_bgr, template_mask, request_dir,
|
||||
instances_raw=template_flower_instances,
|
||||
)
|
||||
visual = flower_count.render_instances(template_bgr, template_count_data)
|
||||
template_sam_visual_file = "flower_count_sam_template.png"
|
||||
cv2.imwrite(os.path.join(request_dir, template_sam_visual_file), visual)
|
||||
except Exception as e:
|
||||
logger.exception("[%s] template flower count failed", request_id)
|
||||
template_sam_error = str(e)
|
||||
|
||||
yolo_visual_file = None
|
||||
if yolo_error is None:
|
||||
try:
|
||||
yolo_overlay = yolo_world.render_boxes(bgr, by_class)
|
||||
yolo_visual_file = "flower_count_yolo.png"
|
||||
cv2.imwrite(os.path.join(request_dir, yolo_visual_file), yolo_overlay)
|
||||
except Exception:
|
||||
logger.exception("[%s] YOLO-World visualization failed", request_id)
|
||||
|
||||
# Vase-identity comparison: needs a matched template name and a vase
|
||||
# mask on both sides (from the batched SAM3 call above). Soft-
|
||||
# skipped (not an error) if either is missing -- e.g. a weak/no
|
||||
# match, or a photo where SAM3 simply didn't find the vase.
|
||||
input_vase_instances = sam3_result.get(("input", config.SAM3_VASE_PROMPT), [])
|
||||
vase_comparison = None
|
||||
if template is not None and sam3_error is not None:
|
||||
vase_comparison = {"error": f"SAM3 unavailable: {sam3_error}"}
|
||||
elif template is not None and not input_vase_instances:
|
||||
vase_comparison = {"error": "No vase detected in your upload."}
|
||||
elif template is not None and template_bgr is not None:
|
||||
try:
|
||||
template_vase_instances = sam3_result.get(
|
||||
("template", config.SAM3_VASE_PROMPT), []
|
||||
)
|
||||
if not template_vase_instances:
|
||||
vase_comparison = {"error": "No vase detected in the matched template photo."}
|
||||
else:
|
||||
input_vase_mask = max(input_vase_instances, key=lambda i: i["score"])["mask"]
|
||||
template_vase_mask = max(template_vase_instances, key=lambda i: i["score"])["mask"]
|
||||
input_vase_crop = vase_compare.crop_mask(bgr, input_vase_mask)
|
||||
template_vase_crop = vase_compare.crop_mask(template_bgr, template_vase_mask)
|
||||
|
||||
result = vase_compare.compare_vases(input_vase_crop, template_vase_crop)
|
||||
input_crop_file = "vase_crop_input.png"
|
||||
template_crop_file = "vase_crop_template.png"
|
||||
cv2.imwrite(os.path.join(request_dir, input_crop_file), input_vase_crop)
|
||||
cv2.imwrite(os.path.join(request_dir, template_crop_file), template_vase_crop)
|
||||
vase_comparison = {
|
||||
**result,
|
||||
"template": template,
|
||||
"input_crop_file": input_crop_file,
|
||||
"template_crop_file": template_crop_file,
|
||||
"error": None,
|
||||
}
|
||||
except Exception as e:
|
||||
logger.exception("[%s] vase comparison failed", request_id)
|
||||
vase_comparison = {"error": str(e)}
|
||||
|
||||
# Flower-count comparison: surfaces a plain-English explanation when
|
||||
# the upload and the matched template simply have different numbers
|
||||
# of flowers -- none of SIFT/ORB/SuperGlue/LoFTR/shape/texture score
|
||||
# "flower count" directly, so a lower match score caused mostly by a
|
||||
# count mismatch would otherwise look unexplained to the user.
|
||||
flower_count_comparison = None
|
||||
if template is not None and template_sam_error is None and template_bgr is not None:
|
||||
flower_count_comparison = _build_flower_count_comparison(
|
||||
count_data.get("total_count", 0), template_count_data.get("total_count", 0), template
|
||||
)
|
||||
elif template is not None and template_sam_error is not None:
|
||||
flower_count_comparison = {"error": f"SAM3 unavailable: {template_sam_error}"}
|
||||
|
||||
yolo_world.unload_model()
|
||||
vase_compare.unload_models()
|
||||
_release_resources()
|
||||
elapsed = time.perf_counter() - start
|
||||
logger.info(
|
||||
"[%s] flower count: %.2fs, SAM3 total=%s, YOLO flower/vase/ribbon=%s/%s/%s, "
|
||||
"vase comparison=%s",
|
||||
request_id, elapsed, count_data.get("total_count"),
|
||||
len(by_class.get("flower", [])), len(by_class.get("vase", [])),
|
||||
len(by_class.get("ribbon", [])) + len(by_class.get("bow", [])),
|
||||
vase_comparison.get("verdict") if vase_comparison and not vase_comparison.get("error") else None,
|
||||
)
|
||||
|
||||
return {
|
||||
"sam": {
|
||||
"total_count": count_data.get("total_count", 0),
|
||||
"clusters": count_data.get("clusters", []),
|
||||
"visual_file": sam_visual_file,
|
||||
"error": sam_error,
|
||||
},
|
||||
"sam_template": {
|
||||
"total_count": template_count_data.get("total_count", 0),
|
||||
"clusters": template_count_data.get("clusters", []),
|
||||
"visual_file": template_sam_visual_file,
|
||||
"error": template_sam_error,
|
||||
},
|
||||
"yolo": {
|
||||
"flower_count": len(by_class.get("flower", [])),
|
||||
"vase_count": len(by_class.get("vase", [])),
|
||||
"ribbon_count": len(by_class.get("ribbon", [])) + len(by_class.get("bow", [])),
|
||||
"visual_file": yolo_visual_file,
|
||||
"error": yolo_error,
|
||||
},
|
||||
"vase_comparison": vase_comparison,
|
||||
"flower_count_comparison": flower_count_comparison,
|
||||
"time_sec": round(elapsed, 3),
|
||||
}
|
||||
|
||||
|
||||
def flower_summary(request_id: str, template: str) -> dict:
|
||||
"""Lightweight companion to the weighted verdict, shown automatically
|
||||
right beside it (not behind the opt-in "Count flowers" button) as soon
|
||||
as a confident match is found: just the flower-count comparison (same
|
||||
message as count_flowers' version) plus a CLIP-only similarity of the
|
||||
flower material itself (SAM3's "flower" masks on both sides, unioned
|
||||
and background-blacked-out, then embedded with CLIP alone).
|
||||
|
||||
Deliberately skips YOLO-World, DINOv2, and the vase comparison -- those
|
||||
stay behind the button since this one runs unconditionally on every
|
||||
confident match and should stay as fast as a SAM3 round trip allows.
|
||||
Still needs the matching pipeline's GPU-resident models unloaded first,
|
||||
same as count_flowers."""
|
||||
request_dir = os.path.join(config.UPLOADS_DIR, request_id)
|
||||
nobg_path = os.path.join(request_dir, "nobg.png")
|
||||
if not os.path.isfile(nobg_path):
|
||||
raise FileNotFoundError(f"no processed upload found for request {request_id}")
|
||||
if template not in _templates_meta:
|
||||
raise ValueError(f"unknown template {template!r}")
|
||||
|
||||
with _pipeline_lock:
|
||||
rgba = cv2.imread(nobg_path, cv2.IMREAD_UNCHANGED)
|
||||
bgr, _, mask = bg_removal.split_rgba(rgba)
|
||||
template_rgba = cv2.imread(
|
||||
_templates_meta[template]["nobg_path"], cv2.IMREAD_UNCHANGED
|
||||
)
|
||||
template_bgr, _, template_mask = bg_removal.split_rgba(template_rgba)
|
||||
|
||||
freed_deep = deep.unload_models()
|
||||
freed_rembg = bg_removal.unload_session()
|
||||
if freed_deep or freed_rembg:
|
||||
logger.info("[%s] unloaded matching-pipeline models before flower summary",
|
||||
request_id)
|
||||
|
||||
count_comparison = None
|
||||
clip_pct = None
|
||||
error = None
|
||||
try:
|
||||
sam3_result = sam3_client.run_jobs(
|
||||
images={"input": bgr, "template": template_bgr},
|
||||
jobs=[
|
||||
{"image": "input", "prompt": config.SAM3_FLOWER_PROMPT,
|
||||
"threshold": config.SAM3_FLOWER_THRESHOLD},
|
||||
{"image": "template", "prompt": config.SAM3_FLOWER_PROMPT,
|
||||
"threshold": config.SAM3_FLOWER_THRESHOLD},
|
||||
],
|
||||
workdir=request_dir,
|
||||
)
|
||||
input_instances = sam3_result.get(("input", config.SAM3_FLOWER_PROMPT), [])
|
||||
template_instances = sam3_result.get(("template", config.SAM3_FLOWER_PROMPT), [])
|
||||
|
||||
input_data = flower_count.count_flowers(bgr, mask, request_dir,
|
||||
instances_raw=input_instances)
|
||||
template_data = flower_count.count_flowers(template_bgr, template_mask, request_dir,
|
||||
instances_raw=template_instances)
|
||||
count_comparison = _build_flower_count_comparison(
|
||||
input_data["total_count"], template_data["total_count"], template
|
||||
)
|
||||
|
||||
input_union = flower_count.union_mask(input_data["_instance_masks"])
|
||||
template_union = flower_count.union_mask(template_data["_instance_masks"])
|
||||
if input_union is not None and template_union is not None:
|
||||
input_crop = vase_compare.crop_mask(bgr, input_union)
|
||||
template_crop = vase_compare.crop_mask(template_bgr, template_union)
|
||||
if input_crop is not None and template_crop is not None:
|
||||
clip_pct = vase_compare.clip_similarity_pct(input_crop, template_crop)
|
||||
except Exception as e:
|
||||
logger.exception("[%s] flower summary failed", request_id)
|
||||
error = str(e)
|
||||
|
||||
vase_compare.unload_models()
|
||||
_release_resources()
|
||||
|
||||
return {
|
||||
"flower_count_comparison": count_comparison,
|
||||
"flower_clip_similarity_pct": clip_pct,
|
||||
"error": error,
|
||||
}
|
||||
|
||||
|
||||
def process_upload(image_bytes: bytes, orig_filename: str) -> dict:
|
||||
bootstrap()
|
||||
_cleanup_old_uploads()
|
||||
|
||||
request_id = uuid.uuid4().hex[:12]
|
||||
logger.info("[%s] new upload: %r (%.1f KB)", request_id, orig_filename,
|
||||
len(image_bytes) / 1024)
|
||||
|
||||
with _pipeline_lock:
|
||||
try:
|
||||
result = _process_upload_locked(request_id, image_bytes, orig_filename)
|
||||
except Exception:
|
||||
logger.exception("[%s] pipeline failed", request_id)
|
||||
raise
|
||||
finally:
|
||||
_release_resources()
|
||||
logger.info("[%s] done, peak RSS so far: %.0f MB", request_id, _peak_rss_mb())
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _process_upload_locked(request_id: str, image_bytes: bytes, orig_filename: str) -> dict:
|
||||
request_dir = os.path.join(config.UPLOADS_DIR, request_id)
|
||||
os.makedirs(request_dir, exist_ok=True)
|
||||
|
||||
total_start = time.perf_counter()
|
||||
|
||||
ext = os.path.splitext(orig_filename)[1].lower() or ".png"
|
||||
original_path = os.path.join(request_dir, "original" + ext)
|
||||
with open(original_path, "wb") as f:
|
||||
f.write(image_bytes)
|
||||
|
||||
bg_start = time.perf_counter()
|
||||
rgba, _ = bg_removal.remove_background_bytes(image_bytes, config.UPLOADS_NOBG_CACHE)
|
||||
bg_elapsed = time.perf_counter() - bg_start
|
||||
logger.info("[%s] background removal: %.2fs", request_id, bg_elapsed)
|
||||
bgr, _, mask = bg_removal.split_rgba(rgba)
|
||||
|
||||
nobg_path = os.path.join(request_dir, "nobg.png")
|
||||
cv2.imwrite(nobg_path, rgba)
|
||||
|
||||
method_results = {}
|
||||
for method in METHODS:
|
||||
results, elapsed, error = _timed(method, _METHOD_RUNNERS[method], bgr, mask)
|
||||
method_results[method] = {"results": results, "time_sec": elapsed, "error": error}
|
||||
logger.info("[%s] %s: %.2fs%s", request_id, method, elapsed,
|
||||
f" (FAILED: {error})" if error else "")
|
||||
# Free this method's GPU allocations before the next one runs rather
|
||||
# than waiting until the whole request finishes -- SuperGlue and
|
||||
# LoFTR are the two that actually use the GPU, and freeing between
|
||||
# them keeps their peak allocations from ever coexisting.
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
for method, payload in method_results.items():
|
||||
results = payload["results"]
|
||||
best = results[0] if results else None
|
||||
payload["best"] = best
|
||||
payload["is_confident"] = bool(
|
||||
best and best["score"] >= config.SCORE_THRESHOLD[method]
|
||||
)
|
||||
if best is not None:
|
||||
template_bgr = cv2.imread(
|
||||
_templates_meta[best["template"]]["nobg_path"], cv2.IMREAD_UNCHANGED
|
||||
)
|
||||
annotated = utils.annotate_score(
|
||||
template_bgr, best["score"], best["confidence"],
|
||||
label=config.METHOD_LABELS[method],
|
||||
)
|
||||
out_name = f"{method}_best.png"
|
||||
cv2.imwrite(os.path.join(request_dir, out_name), annotated)
|
||||
payload["best_image_file"] = out_name
|
||||
|
||||
# Color-space comparison: its own section in the UI (independent of the
|
||||
# method-grid cards above), but also folded into the weighted verdict
|
||||
# below via config.METHOD_WEIGHTS["Color"]. A failure here shouldn't
|
||||
# take down a request that otherwise succeeded.
|
||||
color_start = time.perf_counter()
|
||||
try:
|
||||
color_analysis = color.compare_input_to_templates(bgr, mask)
|
||||
color_error = None
|
||||
except Exception as e:
|
||||
logger.exception("[%s] color analysis failed", request_id)
|
||||
color_analysis = {"input_dominant_colors": [], "templates": []}
|
||||
color_error = str(e)
|
||||
color_elapsed = time.perf_counter() - color_start
|
||||
logger.info("[%s] color analysis: %.2fs%s", request_id, color_elapsed,
|
||||
f" (FAILED: {color_error})" if color_error else "")
|
||||
|
||||
# Shape matching (silhouette/contour, via Hu moments + canonical-aligned
|
||||
# IoU) and texture matching (LBP + GLCM/Haralick) -- both independent
|
||||
# sections, like color space: never fed into weighting or any score.
|
||||
shape_start = time.perf_counter()
|
||||
try:
|
||||
input_shape_data, shape_results = shape_match.compare_to_templates(mask)
|
||||
shape_error = None
|
||||
except Exception as e:
|
||||
logger.exception("[%s] shape analysis failed", request_id)
|
||||
input_shape_data, shape_results = None, []
|
||||
shape_error = str(e)
|
||||
shape_elapsed = time.perf_counter() - shape_start
|
||||
logger.info("[%s] shape analysis: %.2fs%s", request_id, shape_elapsed,
|
||||
f" (FAILED: {shape_error})" if shape_error else "")
|
||||
|
||||
texture_start = time.perf_counter()
|
||||
try:
|
||||
input_texture_data, texture_results = texture_match.compare_to_templates(bgr, mask)
|
||||
texture_error = None
|
||||
except Exception as e:
|
||||
logger.exception("[%s] texture analysis failed", request_id)
|
||||
input_texture_data, texture_results = None, []
|
||||
texture_error = str(e)
|
||||
texture_elapsed = time.perf_counter() - texture_start
|
||||
logger.info("[%s] texture analysis: %.2fs%s", request_id, texture_elapsed,
|
||||
f" (FAILED: {texture_error})" if texture_error else "")
|
||||
|
||||
# Only used for weighting/the returned "methods" dict (so the frontend
|
||||
# can look up Color's raw score for the weighted-breakdown display) --
|
||||
# NOT rendered as a 5th method-grid card, and NOT part of _overall_best.
|
||||
methods_with_color = dict(method_results)
|
||||
methods_with_color["Color"] = _color_as_method_result(color_analysis, color_elapsed, color_error)
|
||||
|
||||
weighted_best, weighted_scores = _weighted_scores(methods_with_color)
|
||||
|
||||
# Color family grid: purely visual, tied to whichever template the
|
||||
# weighted verdict landed on. Skipped if there's no confident match at
|
||||
# all (nothing to compare against). A failure here is likewise soft.
|
||||
family_grid = None
|
||||
if weighted_best is not None:
|
||||
try:
|
||||
family_grid = _build_family_grid(request_dir, bgr, mask, weighted_best)
|
||||
except Exception as e:
|
||||
logger.exception("[%s] family grid failed", request_id)
|
||||
family_grid = {"error": str(e)}
|
||||
|
||||
# Shape/texture visuals: same "tied to whichever template the weighted
|
||||
# verdict landed on" convention as the color family grid.
|
||||
shape_visuals = None
|
||||
if weighted_best is not None and input_shape_data is not None:
|
||||
try:
|
||||
shape_visuals = _build_shape_visuals(request_dir, input_shape_data, weighted_best)
|
||||
except Exception as e:
|
||||
logger.exception("[%s] shape visuals failed", request_id)
|
||||
shape_visuals = None
|
||||
|
||||
texture_visuals = None
|
||||
if weighted_best is not None and input_texture_data is not None:
|
||||
try:
|
||||
texture_visuals = _build_texture_visuals(request_dir, input_texture_data, weighted_best)
|
||||
except Exception as e:
|
||||
logger.exception("[%s] texture visuals failed", request_id)
|
||||
texture_visuals = None
|
||||
|
||||
total_elapsed = time.perf_counter() - total_start
|
||||
logger.info("[%s] total: %.2fs, weighted best: %s", request_id, total_elapsed, weighted_best)
|
||||
|
||||
return {
|
||||
"request_id": request_id,
|
||||
"upload_nobg_file": "nobg.png",
|
||||
"upload_original_file": os.path.basename(original_path),
|
||||
"bg_removal_time_sec": round(bg_elapsed, 3),
|
||||
"total_time_sec": round(total_elapsed, 3),
|
||||
"overall_best": _overall_best(method_results),
|
||||
"method_weights": config.METHOD_WEIGHTS,
|
||||
"weighted_best": weighted_best,
|
||||
"weighted_scores": weighted_scores,
|
||||
"color_analysis": color_analysis,
|
||||
"color_analysis_time_sec": round(color_elapsed, 3),
|
||||
"color_analysis_error": color_error,
|
||||
"family_grid": family_grid,
|
||||
"shape_analysis": {
|
||||
"results": shape_results,
|
||||
"time_sec": round(shape_elapsed, 3),
|
||||
"error": shape_error,
|
||||
"visuals": shape_visuals,
|
||||
},
|
||||
"texture_analysis": {
|
||||
"results": texture_results,
|
||||
"time_sec": round(texture_elapsed, 3),
|
||||
"error": texture_error,
|
||||
"visuals": texture_visuals,
|
||||
},
|
||||
"methods": {
|
||||
method: {
|
||||
"label": config.METHOD_LABELS[method],
|
||||
"time_sec": round(payload["time_sec"], 3),
|
||||
"best": payload["best"],
|
||||
"is_confident": payload["is_confident"],
|
||||
"best_image_file": payload.get("best_image_file"),
|
||||
"results": payload["results"],
|
||||
"error": payload["error"],
|
||||
}
|
||||
for method, payload in methods_with_color.items()
|
||||
},
|
||||
}
|
||||
Reference in New Issue
Block a user