import gradio as gr import cv2 import numpy as np import os import tempfile import time import axengine as axe import common import imgproc rgb_range=255 scale=2 def from_numpy(x): return x if isinstance(x, np.ndarray) else np.array(x) def quantize(img, rgb_range): pixel_range = 255 / rgb_range return np.round(np.clip(img * pixel_range, 0, 255)) / pixel_range # 初始化EDSR和ESPCN模型 def init_SRmodel(EDSR_path="../model_convert/axmodel/edsr_baseline_x2_1.axmodel", ESPCN_path="../model_convert/axmodel/espcn_x2_T9.axmodel"): EDSR_session = axe.InferenceSession(EDSR_path) ESPCN_session = axe.InferenceSession(ESPCN_path) return [EDSR_session, ESPCN_session] SR_sessions=init_SRmodel() def EDSR_infer(frame, EDSR_session=SR_sessions[0]): output_names = [x.name for x in EDSR_session.get_outputs()] input_name = EDSR_session.get_inputs()[0].name lr_y_image, = common.set_channel(frame, n_channels=3) lr_y_image, = common.np_prepare(lr_y_image, rgb_range=rgb_range) sr = EDSR_session.run(output_names, {input_name: lr_y_image}) if isinstance(sr, (list, tuple)): sr = from_numpy(sr[0]) if len(sr) == 1 else [from_numpy(x) for x in sr] else: sr = from_numpy(sr) sr = quantize(sr, rgb_range).squeeze(0) normalized = sr * 255 / rgb_range ndarr = normalized.transpose(1, 2, 0).astype(np.uint8) return ndarr def ESPCN_infer(frame, ESPCN_session=SR_sessions[1]): output_names = [x.name for x in ESPCN_session.get_outputs()] input_name = ESPCN_session.get_inputs()[0].name lr_y_image, lr_cb_image, lr_cr_image = imgproc.preprocess_one_frame(frame) bic_cb_image = cv2.resize(lr_cb_image, (int(lr_cb_image.shape[1] * scale), int(lr_cb_image.shape[0] * scale)), interpolation=cv2.INTER_CUBIC) bic_cr_image = cv2.resize(lr_cr_image, (int(lr_cr_image.shape[1] * scale), int(lr_cr_image.shape[0] * scale)), interpolation=cv2.INTER_CUBIC) sr = ESPCN_session.run(output_names, {input_name: lr_y_image}) if isinstance(sr, (list, tuple)): sr = from_numpy(sr[0]) if len(sr) == 1 else [from_numpy(x) for x in sr] else: sr = from_numpy(sr) ndarr = imgproc.array_to_image(sr) sr_y_image = ndarr.astype(np.float32) / 255.0 sr_ycbcr_image = cv2.merge([sr_y_image[:, :, 0], bic_cb_image, bic_cr_image]) sr_image = imgproc.ycbcr_to_bgr(sr_ycbcr_image) sr_image = np.clip(sr_image* 255.0, 0 , 255).astype(np.uint8) return sr_image # ====================== # 模拟超分辨率模型 # ====================== def EDSR_MODEL(input_data, is_video=False): if is_video: output_frames = [] for frame in input_data: out = EDSR_infer(frame=frame) output_frames.append(out) return output_frames else: out = EDSR_infer(frame=input_data) return out def ESPCN_MODEL(input_data, is_video=False): if is_video: output_frames = [] for frame in input_data: out = ESPCN_infer(frame=frame) output_frames.append(out) return output_frames else: out = ESPCN_infer(frame=input_data) return out # ====================== # 全局状态（单用户） # ====================== class AppState: def __init__(self): self.original_img = None # 原始图（BGR, 高分辨率） self.sr_img = None # 超分图（BGR, 高分辨率） self.is_video = False app_state = AppState() # ====================== # 核心处理函数 # ====================== def process_super_resolution(input_file, model_choice): global app_state if input_file is None: raise gr.Error("请先上传图片或视频！") file_path = input_file app_state = AppState() info_text = "" is_video = any(ext in file_path.lower() for ext in ['.mp4', '.avi', '.mov', '.mkv']) if is_video: # --- 视频处理（直接保存高分辨率）--- cap = cv2.VideoCapture(file_path) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) fps = cap.get(cv2.CAP_PROP_FPS) info_text += f"🎬 视频信息:\n- 总帧数: {total_frames}\n- 帧率: {fps:.2f} FPS\n" frames = [] while True: ret, frame = cap.read() if not ret: break frames.append(frame) cap.release() model_func = EDSR_MODEL if model_choice == "EDSR_MODEL" else ESPCN_MODEL start_time = time.time() output_data = model_func(frames, is_video=True) infer_time = time.time() - start_time info_text += f"\n⏱️ 推理时间: {infer_time:.2f} 秒\n" full_video_path = os.path.join(tempfile.gettempdir(), f"sr_video_x2.mp4") h_out, w_out = output_data[0].shape[:2] info_text += f"- 超分后尺寸: {w_out} x {h_out}\n" fourcc = cv2.VideoWriter_fourcc(*'mp4v') out_video = cv2.VideoWriter(full_video_path, fourcc, fps, (w_out, h_out)) for frame in output_data: out_video.write(frame) out_video.release() app_state.is_video = True return ( gr.update(value=None, visible=False), # image_display gr.update(visible=False), # btn_original gr.update(visible=False), # btn_sr gr.update(value="当前: 无", visible=False), gr.update(value=full_video_path, visible=True), gr.update(value=full_video_path, visible=True), gr.update(visible=False), info_text ) else: # --- 图片处理（保存原始高分辨率）--- img = cv2.imread(file_path) if img is None: raise gr.Error("无法读取图片！") h, w = img.shape[:2] info_text += f"🖼️ 图片信息:\n- 原始尺寸: {w} x {h}\n" app_state.original_img = img.copy() model_func = EDSR_MODEL if model_choice == "EDSR_MODEL" else ESPCN_MODEL start_time = time.time() sr_img = model_func(img, is_video=False) infer_time = time.time() - start_time info_text += f"\n⏱️ 推理时间: {infer_time:.2f} 秒\n" h_out, w_out = sr_img.shape[:2] info_text += f"- 超分后尺寸: {w_out} x {h_out}\n" sr_img_path = os.path.join(tempfile.gettempdir(), f"sr_image_x2.png") cv2.imwrite(sr_img_path, sr_img) app_state.sr_img = sr_img app_state.is_video = False # 默认显示原图（高分辨率，但 UI 会限制尺寸） return ( gr.update(value=app_state.original_img[:, :, ::-1], visible=True), # BGR→RGB gr.update(visible=True), gr.update(visible=True), gr.update(value="当前: 原图", visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(value=sr_img_path, visible=True), info_text ) # ====================== # 切换显示函数（直接使用原始高分辨率图） # ====================== def show_original(): if app_state.original_img is None: return gr.update(), gr.update() # OpenCV BGR → RGB rgb_img = app_state.original_img[:, :, ::-1] return gr.update(value=rgb_img), gr.update(value="当前: 原图") def show_sr(): if app_state.sr_img is None: return gr.update(), gr.update() rgb_img = app_state.sr_img[:, :, ::-1] return gr.update(value=rgb_img), gr.update(value="当前: 超分图") # ====================== # Gradio UI # ====================== with gr.Blocks(title="超分辨率可视化工具", theme=gr.themes.Soft()) as demo: gr.Markdown("## 🚀 超分辨率模型效果可视化") gr.Markdown("上传图片或视频，选择模型，点击箭头切换原图/超分图！") input_file = gr.File( label="📂 上传图片或视频", file_types=["image", "video"], file_count="single" ) with gr.Row(): model_choice = gr.Radio( choices=["EDSR_MODEL", "ESPCN_MODEL"], value="EDSR_MODEL", label="🔍 选择超分辨率模型" ) run_btn = gr.Button("🚀 开始超分", variant="primary") # 图片区：硬性限定尺寸，直接显示原始高分辨率图 with gr.Column(visible=False) as image_section: image_label = gr.Textbox(value="当前: 原图", interactive=False, lines=1) image_display = gr.Image( label="🖼️ 图像显示", width=800, # 👈 固定宽度 height=600 # 👈 固定高度 ) with gr.Row(): btn_original = gr.Button("◀ 原图") btn_sr = gr.Button("超分图 ▶") # 视频区：硬性限定高度 output_video_player = gr.Video( label="▶️ 超分视频（高分辨率）", visible=False, height=450 # 宽度自适应，高度固定 ) with gr.Row(): download_image = gr.File(label="📥 下载超分图片（原图）", visible=False) download_video = gr.File(label="📥 下载超分视频（完整分辨率）", visible=False) info_box = gr.Textbox(label="📊 处理信息", lines=6, interactive=False) run_btn.click( fn=process_super_resolution, inputs=[input_file, model_choice], outputs=[ image_display, btn_original, btn_sr, image_label, output_video_player, download_video, download_image, info_box ] ) btn_original.click(show_original, outputs=[image_display, image_label]) btn_sr.click(show_sr, outputs=[image_display, image_label]) def toggle_ui(file): if file is None: return ( gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False) ) if any(ext in file.lower() for ext in ['.mp4', '.avi', '.mov', '.mkv']): return ( gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True) ) else: return ( gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False) ) input_file.change( fn=toggle_ui, inputs=input_file, outputs=[ image_section, download_image, output_video_player, download_video ] ) if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=7860)