import gradio as gr import numpy as np import cv2 import onnxruntime as ort from PIL import Image, ImageEnhance, ImageFilter, ImageDraw, ImageFont from rembg import remove, new_session import easyocr import torch from spandrel import ModelLoader import os import onnxruntime as ort from transformers import pipeline, AutoModelForImageSegmentation, AutoConfig import json # --- 全局变量与模型缓存 --- models = { "birefnet": None, "bria_session": None, # 必须添加这一行 "ocr": None, "anime_session": None, "upscaler": None } def identify_objects(input_img): global models if input_img is None: return None, json.dumps({"error": "请上传图片"}, ensure_ascii=False) try: if "object_detector" not in models: from transformers import pipeline models["object_detector"] = pipeline( "object-detection", model="hustvl/yolos-tiny", device=-1 ) img_pil = input_img.convert("RGB") results = models["object_detector"](img_pil) draw = ImageDraw.Draw(img_pil) # 准备一个列表存储结构化数据 detections = [] count = 0 for res in results: if res["score"] > 0.5: count += 1 box = res["box"] label = res["label"] score = res["score"] # 绘制逻辑 draw.rectangle([box["xmin"], box["ymin"], box["xmax"], box["ymax"]], outline="red", width=4) # 将信息添加到列表 detections.append({ "id": count, "label": label, "confidence": round(score, 4), # 保留4位小数 "box": { "xmin": int(box["xmin"]), "ymin": int(box["ymin"]), "xmax": int(box["xmax"]), "ymax": int(box["ymax"]) } }) # 构造最终的 JSON 响应 output_json = { "status": "success", "count": count, "detections": detections } # 使用 json.dumps 转换为字符串，ensure_ascii=False 保证中文不乱码 return img_pil, json.dumps(output_json, ensure_ascii=False, indent=4) except Exception as e: error_json = { "status": "error", "message": str(e) } return input_img, json.dumps(error_json, ensure_ascii=False) def process_rembg(input_img): global models if input_img is None: return None try: # 加载模型逻辑 if models["birefnet"] is None: from transformers import AutoModelForImageSegmentation model_id = "ZhengPeng7/BiRefNet_lite" models["birefnet"] = AutoModelForImageSegmentation.from_pretrained( model_id, trust_remote_code=True ) # 补丁：防止某些版本 transformers 报错 if not hasattr(models["birefnet"], "all_tied_weights_keys"): models["birefnet"].all_tied_weights_keys = [] models["birefnet"].to("cpu").eval() # 图像处理 img_pil = input_img.convert("RGB") w, h = img_pil.size # 预处理：1024x1024 归一化 img_resized = img_pil.resize((1024, 1024), Image.BILINEAR) img_np = np.array(img_resized).astype(np.float32) / 255.0 # 标准 ImageNet 归一化参数 mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) std = np.array([0.229, 0.224, 0.225], dtype=np.float32) img_np = (img_np - mean) / std img_np = img_np.transpose((2, 0, 1)) # HWC -> CHW # --- 终极修复：强制 float() --- img_tensor = torch.from_numpy(img_np).unsqueeze(0).float() with torch.no_grad(): # 推理：BiRefNet 返回多尺度输出，取最后一个 preds = models["birefnet"](img_tensor)[-1].sigmoid().cpu() # 还原尺寸 mask = torch.nn.functional.interpolate( preds, size=(h, w), mode='bilinear', align_corners=False )[0, 0] mask_np = (mask.numpy() * 255).astype(np.uint8) # 边缘美化 mask_pil = Image.fromarray(mask_np) mask_pil = mask_pil.filter(ImageFilter.GaussianBlur(radius=0.8)) # 合成 result = img_pil.copy() result.putalpha(mask_pil) return result except Exception as e: import traceback print(traceback.format_exc()) return f"BiRefNet 最终修复版失败: {str(e)}" # 2. 文字识别 (OCR) def process_ocr(input_img): global models if input_img is None: return "未上传图片", "{}" try: # 1. 懒加载 EasyOCR (支持中英文) if "ocr_reader" not in models: print("正在初始化 OCR 阅读器...") models["ocr_reader"] = easyocr.Reader(['ch_sim', 'en'], gpu=False) # 2. 执行识别 img_np = np.array(input_img) results = models["ocr_reader"].readtext(img_np) # 3. 数据处理 full_text = "" # 存储纯文本 ocr_details = [] # 存储 JSON 详情 for (bbox, text, prob) in results: if prob > 0.3: # 累加纯文本 full_text += f"{text}\n" # 整理坐标：bbox 格式为 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]] # 转换为标准的 Python 列表以便 JSON 序列化 box_coords = [[int(x), int(y)] for x, y in bbox] ocr_details.append({ "text": text, "confidence": round(float(prob), 4), "bbox": box_coords }) # 格式化 JSON 字符串 json_output = json.dumps(ocr_details, ensure_ascii=False, indent=2) if not full_text: return "未识别到文字", "{}" return full_text, json_output except Exception as e: return f"OCR 识别出错: {str(e)}", "{}" def process_upscale(input_img): if input_img is None: return None # 更改为 2x 模型路径 model_path = "RealESRGAN_x2plus.pth" if not os.path.exists(model_path): return "请确认 RealESRGAN_x2plus.pth 已上传" try: if models["upscaler"] is None: # 加载模型 loader = ModelLoader() ckpt = loader.load_from_file(model_path) models["upscaler"] = ckpt.model.to("cpu").eval() # 限制线程数，防止 CPU 满载导致网页卡死 torch.set_num_threads(4) # --- 性能预处理 --- img_pil = input_img.convert("RGB") w, h = img_pil.size # 如果原图已经很大（比如超过 1280px），先缩小再 2x 超分 # 这样既保证了细节修复，又保证了速度 max_input = 1024 if max(w, h) > max_input: img_pil.thumbnail((max_input, max_input), Image.LANCZOS) img = np.array(img_pil).astype(np.float32) / 255.0 # HWC -> CHW img_tensor = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0) img_tensor = img_tensor.float() # --- 推理 --- with torch.no_grad(): # 2x 模型的推理速度会比 4x 快非常多 output = models["upscaler"](img_tensor) output = output.squeeze(0).permute(1, 2, 0).clamp(0, 1).numpy() output = (output * 255).astype(np.uint8) result_pil = Image.fromarray(output) # --- 视觉补偿 --- # 因为只放大了 2x，通过锐化来模拟更高倍数的清晰感 enhancer = ImageEnhance.Sharpness(result_pil) result_pil = enhancer.enhance(1.2) return result_pil except Exception as e: import traceback print(traceback.format_exc()) return f"2x 超分失败: {str(e)}" def fast_smart_sharpen(img): """ 通用图像清晰化‘神药’：算法缩放 + 边缘锐化 """ w, h = img.size # 2倍放大 img = img.resize((w * 2, h * 2), Image.LANCZOS) # Unsharp Mask 锐化：能显著增强物体边缘，让模糊的照片变‘硬朗’ img = img.filter(ImageFilter.UnsharpMask(radius=2, percent=120, threshold=2)) # 稍微拉一点对比度，让画面不那么灰 img = ImageEnhance.Contrast(img).enhance(1.05) return img # 4. 动漫化 (使用你上传的 AnimeGANv3_large_Ghibli_c1_e299.onnx) def process_anime(input_img): if input_img is None: return None # 1. 严格检查模型文件 model_path = "AnimeGANv3_large_Ghibli_c1_e299.onnx" if not os.path.exists(model_path): return f"错误：未找到模型文件 {model_path}，请确认已上传到根目录" try: # 2. 确保 Session 存在，如果不存在则立即加载 if models["anime_session"] is None: print(f"正在加载模型: {model_path}...") models["anime_session"] = ort.InferenceSession( model_path, providers=['CPUExecutionProvider'] ) # 3. 获取节点信息 session = models["anime_session"] input_name = session.get_inputs()[0].name output_name = session.get_outputs()[0].name # 4. 预处理 (使用 1024 提升画质) img_pil = input_img.convert("RGB") w, h = img_pil.size process_size = 1024 scale = process_size / max(w, h) if max(w, h) > process_size else 1.0 target_w, target_h = int((w * scale) // 32) * 32, int((h * scale) // 32) * 32 img_resized = img_pil.resize((target_w, target_h), Image.LANCZOS) img_np = np.array(img_resized).astype(np.float32) img_np = (img_np / 127.5) - 1.0 img_np = np.expand_dims(img_np, axis=0) # 5. 执行推理 output = session.run([output_name], {input_name: img_np})[0] # 6. 后处理与画质增强 output = np.squeeze(output) output = (output + 1.0) * 127.5 output = np.clip(output, 0, 255).astype(np.uint8) result_pil = Image.fromarray(output) # 引入增强逻辑 result_pil = ImageEnhance.Sharpness(result_pil).enhance(1.5) result_pil = ImageEnhance.Contrast(result_pil).enhance(1.1) return result_pil.resize((w, h), Image.LANCZOS) except Exception as e: import traceback # 打印详细错误到 Logs，并在 UI 返回简短错误 print(traceback.format_exc()) return f"推理失败: {str(e)}" # --- Gradio 界面设计 --- with gr.Blocks(theme=gr.themes.Default()) as demo: gr.Markdown("# 🚀 ExtractIt AI 多功能平台") with gr.Tabs(): with gr.TabItem("物体识别 (Identify Objects)"): with gr.Row(): input_i = gr.Image(type="pil", label="上传图片") output_t = gr.Textbox(label="识别详情", lines=10, interactive=False) output_i = gr.Image(type="pil", label="可视化结果") btn_i = gr.Button("开始识别") btn_i.click( fn=identify_objects, inputs=input_i, outputs=[output_i, output_t] # 第一个是图片，第二个是文字 ) with gr.TabItem("🖼️ 抠图 (RMBG)"): with gr.Row(): in1 = gr.Image(type="pil", label="上传图片") out1 = gr.Image(label="去背景结果") btn1 = gr.Button("开始处理", variant="primary") btn1.click(process_rembg, inputs=in1, outputs=out1) with gr.TabItem("🎨 动漫化"): with gr.Row(): in2 = gr.Image(type="pil", label="输入照片") out2 = gr.Image(label="二次元化结果") btn2 = gr.Button("风格转换", variant="primary") btn2.click(process_anime, inputs=in2, outputs=out2) with gr.TabItem("🔍 4x 高清修复"): with gr.Row(): in3 = gr.Image(type="pil", label="低清图") out3 = gr.Image(label="4倍超分结果") btn3 = gr.Button("开始增强", variant="primary") btn3.click(process_upscale, inputs=in3, outputs=out3) with gr.TabItem("📝 文字识别 (OCR)"): with gr.Row(): with gr.Column(): input_img = gr.Image(type="pil", label="上传图片") btn_ocr = gr.Button("开始识别文字", variant="primary") with gr.Column(): # 输出识别到的纯文本 output_text = gr.Textbox(label="识别出的文字", lines=8) # 输出 JSON 坐标数据，方便开发者或后续功能使用 output_json = gr.Code(label="文字区块坐标 (JSON)", language="json") btn_ocr.click( fn=process_ocr, inputs=input_img, outputs=[output_text, output_json] ) # 启动队列以支持高并发排队 demo.queue().launch()