"""FastAPI 主服务 — HF Spaces CPU 部署版 改动点： 1. 端口 7860（HF Spaces） 2. MASt3R 已禁用，SceneMatcher 使用 DINOv2-only 3. Qdrant 嵌入式本地文件模式 4. 历史图片管理 API 5. 图片文件服务 API 6. 可视化在删除文件之前生成，支持复杂+文本场景 7. 可视化按需展示（前端点击按钮触发） """ import os import asyncio from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager import time import base64 from fastapi import FastAPI, UploadFile, File, Form, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import FileResponse, JSONResponse from module.model_loader import load_model, load_dinov2, load_ocr, load_bge from module.scene_matcher import SceneMatcher from module.file_manager import ( validate_username, get_user_dir, get_history_images, save_uploaded_file, delete_file ) from module.evaluator import ( evaluate_scene_consistency, evaluate_scene_consistency_cached, update_metadata ) from module.system_monitor import _snapshot, ResourceMonitor, write_session_log from module.qdrant_manager import ( build_user_cache, qdrant_query_scene, has_qdrant_cache, add_to_qdrant, remove_from_qdrant, ) from module.qdrant_manager import qdrant_search_similar, pt_id from module.text_classifier import classify_scene from module.demo_images import ensure_demo_images BASE_DIR = os.path.dirname(os.path.abspath(__file__)) FRONTEND_PATH = os.path.join(BASE_DIR, 'index.html') HISTORY_BASE = os.path.join(BASE_DIR, 'History_imgs') matcher = None dinov2_extractor = None ocr_engine = None bge_tokenizer = None bge_model = None executor = ThreadPoolExecutor(max_workers=2) session_data = { 'startup_stats': None, 'model_loaded_stats': None, 'detections': [], } resource_monitor = ResourceMonitor(sample_interval=0.5) _qd_caches = {} @asynccontextmanager async def lifespan(app): global matcher, dinov2_extractor, ocr_engine, bge_tokenizer, bge_model, session_data startup_t0 = time.perf_counter() session_data['startup_stats'] = _snapshot() model = load_model() dinov2_extractor = load_dinov2() matcher = SceneMatcher(model, dinov2_extractor=dinov2_extractor) ocr_engine = load_ocr() bge_tokenizer, bge_model = load_bge() session_data['model_loaded_stats'] = _snapshot() model_load_time = time.perf_counter() - startup_t0 print(f'[启动] 模型加载完成: {model_load_time:.1f}s') qdrant_t0 = time.perf_counter() if os.path.isdir(HISTORY_BASE): for uname in os.listdir(HISTORY_BASE): udir = os.path.join(HISTORY_BASE, uname) if not os.path.isdir(udir): continue imgs = get_history_images(udir) if not imgs: continue try: if not has_qdrant_cache(uname): print(f"[qdrant] 启动预构建: 用户 [{uname}]") build_user_cache(uname, udir, ocr_engine, dinov2_extractor, bge_tokenizer, bge_model) _qd_caches[uname] = { 'text': qdrant_query_scene(uname, 'text'), 'complex': qdrant_query_scene(uname, 'complex'), } _t_n = len(_qd_caches[uname]['text']) _c_n = len(_qd_caches[uname]['complex']) print(f" [{uname}] 缓存已就绪 (text={_t_n}, complex={_c_n})") except Exception as e: print(f" [{uname}] 缓存加载失败: {e}") qdrant_load_time = time.perf_counter() - qdrant_t0 print(f"[qdrant] 缓存检查完成: {len(_qd_caches)} 个用户 ({qdrant_load_time:.1f}s)") total_startup = time.perf_counter() - startup_t0 print(f'[启动] 总启动时间: {total_startup:.1f}s') # 生成示例场景图片（用于 /api/demo 端点） ensure_demo_images(BASE_DIR) yield session_data['shutdown_stats'] = _snapshot() try: log_path = write_session_log(session_data) print(f"日志已写入: {log_path}") except Exception as e: print(f"日志写入失败: {e}") app = FastAPI(lifespan=lifespan) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) @app.get("/") async def serve_frontend(): return FileResponse(FRONTEND_PATH) @app.get("/cover.png") async def serve_cover(): """提供背景图片 cover.png""" cover_path = os.path.join(BASE_DIR, 'cover.png') if os.path.exists(cover_path): return FileResponse(cover_path, media_type='image/png') raise HTTPException(status_code=404, detail="cover not found") @app.post("/api/login") async def login(username: str = Form(...)): try: uname = validate_username(username) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) user_dir = get_user_dir(uname) history = get_history_images(user_dir) cached = uname in _qd_caches if not cached and history and has_qdrant_cache(uname): _qd_caches[uname] = { 'text': qdrant_query_scene(uname, 'text'), 'complex': qdrant_query_scene(uname, 'complex'), } cached = True return { "username": uname, "history_count": len(history), "cached": cached, } # ==================== 历史图片管理 API ==================== @app.get("/api/history") async def list_history(username: str): try: uname = validate_username(username) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) user_dir = get_user_dir(uname) images = get_history_images(user_dir) result = [] for img_path in images: fname = os.path.basename(img_path) fsize = os.path.getsize(img_path) if os.path.exists(img_path) else 0 result.append({ 'filename': fname, 'size_kb': round(fsize / 1024, 1), 'url': f'/api/image/{uname}/{fname}', }) return {"username": uname, "images": result, "count": len(result)} @app.get("/api/all_users") async def list_all_users(): users = [] if os.path.isdir(HISTORY_BASE): for uname in sorted(os.listdir(HISTORY_BASE)): udir = os.path.join(HISTORY_BASE, uname) if not os.path.isdir(udir): continue imgs = get_history_images(udir) if imgs: users.append({ 'username': uname, 'image_count': len(imgs), }) return {"users": users} @app.get("/api/image/{username}/{filename}") async def serve_image(username: str, filename: str): uname = validate_username(username) safe_name = os.path.basename(filename) user_dir = get_user_dir(uname) img_path = os.path.join(user_dir, safe_name) if not os.path.exists(img_path): raise HTTPException(status_code=404, detail="图片不存在") if not os.path.abspath(img_path).startswith(os.path.abspath(user_dir)): raise HTTPException(status_code=403, detail="无权访问") return FileResponse(img_path) @app.post("/api/delete_image") async def delete_image(username: str = Form(...), filename: str = Form(...)): global _qd_caches try: uname = validate_username(username) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) safe_name = os.path.basename(filename) user_dir = get_user_dir(uname) img_path = os.path.join(user_dir, safe_name) if not os.path.exists(img_path): raise HTTPException(status_code=404, detail="图片不存在") if not os.path.abspath(img_path).startswith(os.path.abspath(user_dir)): raise HTTPException(status_code=403, detail="无权操作") try: remove_from_qdrant(uname, img_path, scene_type=None) if uname in _qd_caches: _qd_caches[uname] = { 'text': qdrant_query_scene(uname, 'text'), 'complex': qdrant_query_scene(uname, 'complex'), } except Exception as e: print(f" [删除] Qdrant 清理失败: {e}") delete_file(img_path) from module.file_manager import load_metadata, save_metadata metadata = load_metadata(user_dir) metadata['images'] = [m for m in metadata['images'] if m.get('filename') != safe_name] save_metadata(user_dir, metadata) remaining = get_history_images(user_dir) return { "success": True, "deleted": safe_name, "remaining_count": len(remaining), } # ==================== 检测 API ==================== def _run_detect_pipeline(uploaded_paths, uname, ocr_engine, bge_tokenizer, bge_model, dinov2_extractor, matcher, history_before, scene_mode='auto'): scene_type_q = 'complex' ocr_result_q = None precomputed_text = None precomputed_bge = None precomputed_dinov2 = None _t0 = time.perf_counter() if scene_mode in ('text', 'complex'): # 用户手动指定场景，跳过 classify_scene scene_type_q = scene_mode print(f' [SKIP] classify_scene (用户指定场景: {scene_mode})') if scene_mode == 'text' and ocr_engine is not None and uploaded_paths: # 文本场景仍需 OCR 提取文本 _, _, detail, ocr_result_q = classify_scene(uploaded_paths[0], ocr_engine) precomputed_text = detail.get('full_text', '') if precomputed_text and precomputed_text.strip(): from module.qdrant_manager import _bge_encode precomputed_bge = _bge_encode(precomputed_text, bge_tokenizer, bge_model) _t1 = time.perf_counter() print(f' [TIME] ocr+bge (skip classify): {_t1-_t0:.3f}s') elif ocr_engine is not None and uploaded_paths: scene_type_q, _, _, ocr_result_q = classify_scene(uploaded_paths[0], ocr_engine) _t1 = time.perf_counter() print(f' [TIME] classify_scene: {_t1-_t0:.3f}s') if ocr_result_q is not None and ocr_result_q and ocr_result_q[0]: precomputed_text = ' '.join(line[1][0] for line in ocr_result_q[0]) if precomputed_text and precomputed_text.strip(): from module.qdrant_manager import _bge_encode precomputed_bge = _bge_encode(precomputed_text, bge_tokenizer, bge_model) _t2 = time.perf_counter() print(f' [TIME] bge_encode: {_t2-_t1:.3f}s') user_qd = _qd_caches.get(uname, {}) has_cache = isinstance(user_qd, dict) and bool(user_qd) cache_used = False final_result = True evaluation_logs = [] best_match = None if has_cache: cache_used = True final_result, evaluation_logs, best_match = evaluate_scene_consistency_cached( matcher, uploaded_paths, uname, ocr_engine, bge_tokenizer, bge_model, scene_type_q, qdrant_search_similar, dinov2_extractor, precomputed_ocr_result=ocr_result_q, precomputed_bge_vec=precomputed_bge, qd_cache=user_qd, ) if not cache_used: final_result, evaluation_logs, best_match = evaluate_scene_consistency( matcher, uploaded_paths, history_before, ocr_engine, bge_tokenizer, bge_model, ) same_scene_paths = set() for log in evaluation_logs: if log['is_same_scene'] and log['query_path'] in uploaded_paths: same_scene_paths.add(log['query_path']) # ★ 先把需要可视化的图片数据读入内存，再删除文件 ★ # 这样可视化构建不依赖磁盘文件，且不会被删除操作影响 _image_cache = {} for log in evaluation_logs: if not log.get('is_same_scene'): continue for key in ('query_path', 'history_path'): p = log.get(key, '') if p and p not in _image_cache and os.path.exists(p): try: import cv2, numpy as np data = np.fromfile(p, dtype=np.uint8) img = cv2.imdecode(data, cv2.IMREAD_COLOR) if img is not None: _image_cache[p] = img except Exception: pass # 删除同一场景的文件 for path in same_scene_paths: delete_file(path) # 构建可视化（使用内存中的图片数据） visualization_data = [] for log in evaluation_logs: if not log.get('is_same_scene'): continue scene_type = log.get('scene_type', 'complex') try: if scene_type == 'complex': vis = _build_complex_vis_from_cache(log, _image_cache) if vis: visualization_data.append(vis) elif scene_type == 'text': vis = _build_text_vis_from_cache(log, _image_cache) if vis: visualization_data.append(vis) except Exception as e: print(f" [可视化] 生成失败: {e}") per_image_results = {} for orig, path in [(os.path.basename(p), p) for p in uploaded_paths]: if path in same_scene_paths: per_image_results[path] = '未通过（同一场景，已删除）' else: per_image_results[path] = '通过' if final_result else '未通过' return final_result, evaluation_logs, best_match, scene_type_q, same_scene_paths, per_image_results, cache_used, precomputed_text, precomputed_bge, precomputed_dinov2, visualization_data @app.post("/api/detect") async def detect( username: str = Form(...), scene_mode: str = Form(default='auto'), images: list[UploadFile] = File(default=[]), ): global session_data, _qd_caches try: uname = validate_username(username) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) user_dir = get_user_dir(uname) if len(images) == 0: raise HTTPException(status_code=400, detail="请上传至少一张图片") history_before = get_history_images(user_dir) saved_info = [] for f in images: data = await f.read() if not data: continue path = save_uploaded_file(data, f.filename, user_dir) saved_info.append((f.filename, path)) if not saved_info: raise HTTPException(status_code=400, detail="没有有效的上传图片") uploaded_paths = [p for _, p in saved_info] if len(history_before) == 0: # 清理可能的 Qdrant 残留数据（用户可能通过历史管理删除了所有图片） try: from module.qdrant_manager import qdrant_delete_user_collections qdrant_delete_user_collections(uname) except Exception: pass _qd_caches.pop(uname, None) per_image_results = {p: '通过（首次上传，跳过比对）' for _, p in saved_info} update_metadata(user_dir, saved_info, per_image_results) try: for _, p in saved_info: add_to_qdrant(uname, p, ocr_engine, dinov2_extractor, bge_tokenizer, bge_model) _qd_caches[uname] = { 'text': qdrant_query_scene(uname, 'text'), 'complex': qdrant_query_scene(uname, 'complex'), } except Exception: pass return { "result": "pass", "reason": "first_upload", "saved_files": [ {"filename": os.path.basename(p), "original_name": orig, "deleted": False} for orig, p in saved_info ], "evaluation_logs": [], "best_match": None, "visualizations": [], } pre_stats = _snapshot() detect_start_time = asyncio.get_event_loop().time() resource_monitor.start_monitoring() loop = asyncio.get_event_loop() final_result, evaluation_logs, best_match, scene_type_q, same_scene_paths, per_image_results, cache_used, precomp_text, precomp_bge, precomp_dino, visualization_data = \ await loop.run_in_executor( executor, _run_detect_pipeline, uploaded_paths, uname, ocr_engine, bge_tokenizer, bge_model, dinov2_extractor, matcher, history_before, scene_mode, ) peak = resource_monitor.stop_monitoring() detect_duration = asyncio.get_event_loop().time() - detect_start_time kept_saved_info = [(orig, p) for orig, p in saved_info if p not in same_scene_paths] update_metadata(user_dir, kept_saved_info, per_image_results) kept_paths_qd = [p for _, p in kept_saved_info] removed_paths_qd = list(same_scene_paths) st_q = scene_type_q def _post_update_qdrant(): for path in kept_paths_qd: try: add_to_qdrant(uname, path, ocr_engine, dinov2_extractor, bge_tokenizer, bge_model, scene_type=st_q, precomputed_text=precomp_text, precomputed_bge=precomp_bge, precomputed_dinov2=precomp_dino) _qd_caches[uname] = { 'text': qdrant_query_scene(uname, 'text'), 'complex': qdrant_query_scene(uname, 'complex'), } except Exception: pass for path in removed_paths_qd: try: remove_from_qdrant(uname, path, scene_type=None) except Exception: pass loop.run_in_executor(executor, _post_update_qdrant) best_match_response = None if best_match: best_match_response = { 'query_image': best_match['query_image'], 'history_image': best_match['history_image'], 'similarity_score': best_match['similarity_score'], 'match_count': best_match['match_count'], 'inlier_ratio': best_match['inlier_ratio'], 'avg_confidence': best_match['avg_confidence'], } return { "result": "pass" if final_result else "fail", "saved_files": [ {"filename": os.path.basename(p), "original_name": orig, "deleted": p in same_scene_paths} for orig, p in saved_info ], "evaluation_logs": [ { "query_image": log['query_image'], "history_image": log['history_image'], "scene_type": log.get('scene_type', 'complex'), "similarity_score": log['similarity_score'], "match_count": log['match_count'], "dinov2_similarity": log.get('dinov2_similarity'), "text_similarity": log.get('text_similarity'), "bge_search_score": log.get('bge_search_score'), "is_same_scene": log['is_same_scene'], "error": log.get('error'), } for log in evaluation_logs ], "best_match": best_match_response, "visualizations": visualization_data, } def _build_complex_vis_from_cache(log_entry, image_cache): """复杂场景可视化：DINOv2 patch 匹配连线 图片数据从 image_cache 中获取（内存），不依赖磁盘文件。 """ try: import cv2 import numpy as np q_path = log_entry.get('query_path', '') h_path = log_entry.get('history_path', '') patch_info = log_entry.get('patch_match_info') img1 = image_cache.get(q_path) img2 = image_cache.get(h_path) if img1 is None or img2 is None: return None if patch_info is None: return None orig_h1, orig_w1 = img1.shape[:2] orig_h2, orig_w2 = img2.shape[:2] max_h = 600 scale1 = max_h / orig_h1 scale2 = max_h / orig_h2 img1_disp = cv2.resize(img1, (int(orig_w1*scale1), max_h)) img2_disp = cv2.resize(img2, (int(orig_w2*scale2), max_h)) h1, w1 = img1_disp.shape[:2] h2, w2 = img2_disp.shape[:2] matches = patch_info.get('matches', []) q_grid = patch_info.get('query_grid', (37, 37)) h_grid = patch_info.get('hist_grid', (37, 37)) q_orig_size = patch_info.get('query_image_size', (orig_w1, orig_h1)) h_orig_size = patch_info.get('hist_image_size', (orig_w2, orig_h2)) q_crop = patch_info.get('query_crop_offset', (0, 0)) h_crop = patch_info.get('hist_crop_offset', (0, 0)) q_resize = patch_info.get('query_resize_size', (518, 518)) h_resize = patch_info.get('hist_resize_size', (518, 518)) q_n_h, q_n_w = q_grid h_n_h, h_n_w = h_grid patch_size = 14 q_scale_resize = q_resize[1] / q_orig_size[1] h_scale_resize = h_resize[1] / h_orig_size[1] gap = 20 canvas = np.ones((max_h, w1 + gap + w2, 3), dtype=np.uint8) * 240 canvas[:h1, :w1] = img1_disp canvas[:h2, w1+gap:] = img2_disp def _patch_center(idx, n_w, crop_off, scale_r, disp_s, is_right, gap_off): row, col = idx // n_w, idx % n_w cx = int(((col + 0.5) * patch_size + crop_off[0]) / scale_r * disp_s) cy = int(((row + 0.5) * patch_size + crop_off[1]) / scale_r * disp_s) if is_right: cx += gap_off return cx, cy for qi, hi, sim in matches: q_cx, q_cy = _patch_center(qi, q_n_w, q_crop, q_scale_resize, scale1, False, 0) h_cx, h_cy = _patch_center(hi, h_n_w, h_crop, h_scale_resize, scale2, True, w1 + gap) intensity = min(1.0, max(0.0, (sim - 0.3) / 0.7)) color = (0, int(200 * intensity + 55), int(255 * (1 - intensity))) cv2.line(canvas, (q_cx, q_cy), (h_cx, h_cy), color, 1, cv2.LINE_AA) n_matches = len(matches) avg_sim = sum(s for _, _, s in matches) / max(n_matches, 1) _, buffer = cv2.imencode('.jpg', canvas, [cv2.IMWRITE_JPEG_QUALITY, 95]) img_b64 = base64.b64encode(buffer).decode('utf-8') return { 'scene_type': 'complex', 'query_image': log_entry['query_image'], 'history_image': log_entry['history_image'], 'dinov2_similarity': patch_info.get('cls_similarity', 0), 'patch_match_count': n_matches, 'avg_patch_similarity': round(avg_sim, 4), 'image_base64': f'data:image/jpeg;base64,{img_b64}', } except Exception as e: print(f" [可视化-复杂] 生成失败: {e}") return None def _build_text_vis_from_cache(log_entry, image_cache): """文本场景可视化：两张图并排 + 标注 图片数据从 image_cache 中获取（内存），不依赖磁盘文件。 """ try: import cv2 import numpy as np from PIL import Image, ImageDraw, ImageFont import re q_path = log_entry.get('query_path', '') h_path = log_entry.get('history_path', '') img1 = image_cache.get(q_path) img2 = image_cache.get(h_path) if img1 is None or img2 is None: return None text1 = log_entry.get('query_text', '') text2 = log_entry.get('history_text', '') text_sim = log_entry.get('text_similarity', 0) words1 = set(re.findall(r'[a-zA-Z\u4e00-\u9fff]{2,}', text1)) words2 = set(re.findall(r'[a-zA-Z\u4e00-\u9fff]{2,}', text2)) common_words = words1 & words2 diff_words1 = words1 - words2 diff_words2 = words2 - words1 max_h = 600 scale1 = max_h / img1.shape[0] scale2 = max_h / img2.shape[0] img1_s = cv2.resize(img1, (int(img1.shape[1]*scale1), max_h)) img2_s = cv2.resize(img2, (int(img2.shape[1]*scale2), max_h)) h1, w1 = img1_s.shape[:2] h2, w2 = img2_s.shape[:2] gap = 10 total_w = w1 + gap + w2 canvas = np.ones((max_h, total_w, 3), dtype=np.uint8) * 240 canvas[:max_h, :w1] = img1_s canvas[:max_h, w1+gap:w1+gap+w2] = img2_s canvas_rgb = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB) pil_img = Image.fromarray(canvas_rgb) draw = ImageDraw.Draw(pil_img) try: font_large = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 16) except Exception: try: font_large = ImageFont.truetype("DejaVuSans.ttf", 16) except Exception: font_large = ImageFont.load_default() draw.text((10, 8), "Query (OCR text)", fill=(0, 0, 0), font=font_large) draw.text((w1+gap+10, 8), "History (OCR text)", fill=(0, 0, 0), font=font_large) common_list = sorted(common_words)[:20] canvas = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR) _, buffer = cv2.imencode('.jpg', canvas, [cv2.IMWRITE_JPEG_QUALITY, 95]) img_b64 = base64.b64encode(buffer).decode('utf-8') return { 'scene_type': 'text', 'query_image': log_entry['query_image'], 'history_image': log_entry['history_image'], 'text_similarity': round(text_sim, 4), 'common_keyword_count': len(common_words), 'diff_keyword_count_query': len(diff_words1), 'diff_keyword_count_history': len(diff_words2), 'common_keywords': common_list, 'image_base64': f'data:image/jpeg;base64,{img_b64}', } except Exception as e: print(f" [可视化-文本] 生成失败: {e}") return None @app.get("/api/status") async def status(): return { "model_loaded": matcher is not None, "device": matcher.device if matcher else None, "ocr_available": ocr_engine is not None, "bge_available": bge_tokenizer is not None, "qdrant_cached_users": len(_qd_caches), } @app.post("/api/demo") async def demo(scene_type: str = Form("auto")): """示例场景检测：自动读取 test_imgs 中的示例图片执行检测 Args: scene_type: 'text' / 'complex' / 'auto'（auto 表示自动选择） """ import shutil demo_map = ensure_demo_images(BASE_DIR) if not demo_map: return JSONResponse({"error": "示例图片生成失败"}, status_code=500) # 选择场景类型 if scene_type == 'auto': scene_type = 'text' if scene_type not in demo_map: return JSONResponse({"error": f"不支持的场景类型: {scene_type}"}, status_code=400) img_paths = demo_map[scene_type] if len(img_paths) < 2: return JSONResponse({"error": "示例图片不足"}, status_code=500) # 复制到临时目录执行检测（不污染任何用户的历史数据） import tempfile tmp_dir = tempfile.mkdtemp(prefix='demo_') tmp_imgs = [] for i, p in enumerate(img_paths[:2]): dst = os.path.join(tmp_dir, f"demo_{i}.jpg") shutil.copy(p, dst) tmp_imgs.append(dst) try: loop = asyncio.get_event_loop() result = await loop.run_in_executor( executor, _run_demo_pipeline, tmp_imgs, scene_type ) return result except Exception as e: import traceback traceback.print_exc() return JSONResponse({"error": f"示例检测失败: {str(e)}"}, status_code=500) finally: # 清理临时目录 try: shutil.rmtree(tmp_dir, ignore_errors=True) except Exception: pass def _run_demo_pipeline(img_paths, expected_scene): """执行示例检测流程（不涉及用户历史数据，纯两张图比对） expected_scene 直接指定场景类型，跳过 classify_scene 节省时间 """ import cv2 import numpy as np path1, path2 = img_paths[0], img_paths[1] # 跳过场景分类，直接使用 expected_scene scene_type = expected_scene ocr_result = None full_text = '' print(f"[demo] 跳过场景分类，直接使用指定场景: {scene_type}") is_same = False similarity = 0.0 log_entry = { 'query_path': path1, 'history_path': path2, 'scene_type': scene_type, 'is_same_scene': False, 'query_image': os.path.basename(path1), 'history_image': os.path.basename(path2), 'query_text': full_text, } visualization_data = [] if scene_type == 'text': # 文本场景：OCR + BGE 比对 from module.text_matcher import compute_text_similarity try: # 直接 OCR 提取两张图文本，跳过 classify_scene from module.text_classifier import _ocr_image full_text = _ocr_image(path1, ocr_engine) or '' text2 = _ocr_image(path2, ocr_engine) or '' log_entry['query_text'] = full_text log_entry['history_text'] = text2 # BGE 编码比对 text_sim = compute_text_similarity(full_text, text2, bge_tokenizer, bge_model) similarity = text_sim is_same = text_sim >= 0.85 log_entry['text_similarity'] = text_sim log_entry['is_same_scene'] = is_same except Exception as e: print(f"[demo] 文本场景比对失败: {e}") else: # 复杂场景：DINOv2 比对 try: if dinov2_extractor is not None: dino_sim = dinov2_extractor.compute_similarity(path1, path2) similarity = dino_sim if dino_sim else 0.0 is_same = similarity >= 0.5 log_entry['dinov2_similarity'] = round(similarity, 4) log_entry['is_same_scene'] = is_same # 计算 patch 匹配（用于可视化） if is_same: patch_info = dinov2_extractor.compute_patch_matches(path1, path2, top_k=50) log_entry['patch_match_info'] = patch_info except Exception as e: print(f"[demo] 复杂场景比对失败: {e}") # 构建可视化（两张图都还在，直接构建） _image_cache = {} for p in [path1, path2]: try: data = np.fromfile(p, dtype=np.uint8) img = cv2.imdecode(data, cv2.IMREAD_COLOR) if img is not None: _image_cache[p] = img except Exception: pass try: if scene_type == 'text': vis = _build_text_vis_from_cache(log_entry, _image_cache) if vis: visualization_data.append(vis) else: vis = _build_complex_vis_from_cache(log_entry, _image_cache) if vis: visualization_data.append(vis) except Exception as e: print(f"[demo] 可视化生成失败: {e}") # 将原图也转为 base64 返回（前端无需再请求 /api/image） def _img_to_b64(path): try: data = np.fromfile(path, dtype=np.uint8) img = cv2.imdecode(data, cv2.IMREAD_COLOR) if img is None: return None _, buf = cv2.imencode('.jpg', img, [cv2.IMWRITE_JPEG_QUALITY, 95]) return f'data:image/jpeg;base64,{base64.b64encode(buf).decode("utf-8")}' except Exception: return None return { "scene_type": scene_type, "is_same_scene": is_same, "similarity": round(similarity, 4), "query_image": os.path.basename(path1), "history_image": os.path.basename(path2), "query_image_b64": _img_to_b64(path1), "history_image_b64": _img_to_b64(path2), "visualizations": visualization_data, "evaluation_logs": [{ "query_image": os.path.basename(path1), "history_image": os.path.basename(path2), "scene_type": scene_type, "is_same_scene": is_same, "similarity": round(similarity, 4), }], } if __name__ == "__main__": import uvicorn port = int(os.environ.get("PORT", 7860)) uvicorn.run("app:app", host="0.0.0.0", port=port, reload=False)