choose source and arrow for generated videos

2026-05-03 02:51:27 +08:00
parent 2fcab9c71a
commit 5ba2d48fdb
3 changed files with 70 additions and 16 deletions
--- a/WebSite/src/App.tsx
+++ b/WebSite/src/App.tsx
@@ -85,6 +85,13 @@ const DEFAULT_PACKAGE_OPTIONS: PackageOptions = {
  images: IMAGE_PACKAGE_OPTIONS.map(option => option.key),
 };

+const VIDEO_SOURCE_OPTIONS = [
+  { key: 'original', label: '原始序列' },
+  { key: 'hard_boundary', label: '硬边界' },
+  { key: 'gaussian_smooth', label: '高斯平滑' },
+  { key: 'soft_transition', label: '软过渡重建' },
+];
+
 type LibraryItem = {
  id: string;
  patientId: string;
@@ -292,6 +299,8 @@ export default function App() {
  const [packageOptions, setPackageOptions] = useState<PackageOptions>(DEFAULT_PACKAGE_OPTIONS);
  const [videoMaxAngle, setVideoMaxAngle] = useState(20);
  const [videoDuration, setVideoDuration] = useState(6);
+  const [videoSource, setVideoSource] = useState('original');
+  const [showVideoArrow, setShowVideoArrow] = useState(true);

  // --- User Management Shared State ---
  const [newUsername, setNewUsername] = useState('');
@@ -303,6 +312,11 @@ export default function App() {

  const selectedDataset = libraryData.find(item => item.id === selectedLibraryId) || libraryData[0];
  const selectedInputDir = selectedDataset?.dicomPath || '';
+  const selectedVideoSource = VIDEO_SOURCE_OPTIONS.find(option => option.key === videoSource) || VIDEO_SOURCE_OPTIONS[0];
+  const videoSourceInputDir = videoSource === 'original'
+    ? selectedInputDir
+    : deformationJob?.result?.outputs?.[videoSource] || '';
+  const isVideoSourceReady = Boolean(videoSourceInputDir);

  useEffect(() => {
    if (!activeUserMenu) return;
@@ -793,19 +807,24 @@ export default function App() {

  const handleGenerateVideo = async () => {
    if (videoJob?.status === 'running') return;
+    if (!videoSourceInputDir) {
+      showToast(videoSource === 'original' ? '请选择影像库数据源' : '请先完成四状态输出，再生成该状态视频');
+      return;
+    }
    try {
      const job = await apiRequest('/api/video', {
        method: 'POST',
        body: JSON.stringify({
-          inputDir: selectedInputDir,
+          inputDir: videoSourceInputDir,
          maxAngle: videoMaxAngle,
-          durationSeconds: videoDuration
+          durationSeconds: videoDuration,
+          showArrow: showVideoArrow,
        })
      }) as BackendJob;
      setVideoJob(job);
      setBackendOnline(true);
-      setBackendMessage('generate_head_extension_video.py 任务已提交');
-      showToast('视频任务已提交');
+      setBackendMessage(`generate_head_extension_video.py / ${selectedVideoSource.label} 任务已提交`);
+      showToast(`${selectedVideoSource.label} 视频任务已提交`);
    } catch (error) {
      setBackendOnline(false);
      showToast((error as Error).message);
@@ -1029,6 +1048,24 @@ export default function App() {
                    <label className="text-[10px] font-bold text-slate-400 uppercase tracking-widest">视频生成</label>
                    <Film size={16} className="text-blue-600" />
                  </div>
+                  <div>
+                    <div className="flex justify-between text-[10px] font-bold mb-2 text-slate-500">
+                      <span>视频来源序列</span>
+                      {videoSource !== 'original' && !isVideoSourceReady && <span className="text-amber-500">需先生成四状态</span>}
+                    </div>
+                    <select
+                      value={videoSource}
+                      onChange={event => {
+                        setVideoSource(event.target.value);
+                        setVideoJob(null);
+                      }}
+                      className="w-full px-3 py-2.5 bg-slate-50 border border-slate-200 rounded-xl outline-none focus:ring-2 focus:ring-blue-500 text-xs font-bold text-slate-700"
+                    >
+                      {VIDEO_SOURCE_OPTIONS.map(option => (
+                        <option key={option.key} value={option.key}>{option.label}</option>
+                      ))}
+                    </select>
+                  </div>
                  <div className="grid grid-cols-2 gap-4">
                    <div>
                      <div className="flex justify-between text-[10px] font-bold mb-2 text-slate-500"><span>最大角度</span><span>{videoMaxAngle}°</span></div>
@@ -1039,8 +1076,19 @@ export default function App() {
                      <input type="range" min="3" max="12" step="1" value={videoDuration} onChange={e => setVideoDuration(parseInt(e.target.value, 10))} className="w-full accent-blue-600" />
                    </div>
                  </div>
-                  <button onClick={handleGenerateVideo} disabled={videoJob?.status === 'running' || !selectedInputDir} className="w-full py-3 bg-slate-100 text-slate-700 rounded-xl text-xs font-bold hover:bg-blue-600 hover:text-white transition-all flex items-center justify-center gap-2 disabled:opacity-50">
-                    <RefreshCw size={14} className={videoJob?.status === 'running' ? 'animate-spin' : ''} /> {videoJob?.status === 'running' ? '正在生成视频' : '生成角度变化视频'}
+                  <button
+                    onClick={() => setShowVideoArrow(value => !value)}
+                    className={`w-full py-2.5 rounded-xl text-xs font-black transition-all flex items-center justify-center gap-2 border ${
+                      showVideoArrow
+                        ? 'bg-yellow-50 text-yellow-700 border-yellow-200 hover:bg-yellow-100'
+                        : 'bg-slate-50 text-slate-500 border-slate-200 hover:bg-slate-100'
+                    }`}
+                  >
+                    {showVideoArrow ? <Eye size={14} /> : <EyeOff size={14} />}
+                    {showVideoArrow ? '隐藏视频黄色箭头' : '显示视频黄色箭头'}
+                  </button>
+                  <button onClick={handleGenerateVideo} disabled={videoJob?.status === 'running' || !isVideoSourceReady} className="w-full py-3 bg-slate-100 text-slate-700 rounded-xl text-xs font-bold hover:bg-blue-600 hover:text-white transition-all flex items-center justify-center gap-2 disabled:opacity-50">
+                    <RefreshCw size={14} className={videoJob?.status === 'running' ? 'animate-spin' : ''} /> {videoJob?.status === 'running' ? '正在生成视频' : `生成${selectedVideoSource.label}视频`}
                  </button>
                  {videoJob?.status === 'completed' && videoJob.result?.video?.path ? (
                    <div className="space-y-3">
--- a/generate_head_extension_video.py
+++ b/generate_head_extension_video.py
@@ -74,7 +74,7 @@ def smoothstep(t):
    return t * t * (3 - 2 * t)


-def make_frame(before_image, angle, max_angle):
+def make_frame(before_image, angle, max_angle, show_arrow=True):
    after_image = video_soft_bend_2d(before_image, angle)

    frame = Image.new("RGB", (1920, 1080), (0, 0, 0))
@@ -95,11 +95,11 @@ def make_frame(before_image, angle, max_angle):
        fill=(255, 255, 255),
    )

-    # Yellow direction arrow on the animated side.
    arrow = (255, 210, 60)
-    x0, y0, x1, y1 = 1390, 405, 1515, 335
-    draw.line((x0, y0, x1, y1), fill=arrow, width=8)
-    draw.polygon([(x1, y1), (x1 - 34, y1 + 3), (x1 - 12, y1 + 29)], fill=arrow)
+    if show_arrow:
+        x0, y0, x1, y1 = 1390, 405, 1515, 335
+        draw.line((x0, y0, x1, y1), fill=arrow, width=8)
+        draw.polygon([(x1, y1), (x1 - 34, y1 + 3), (x1 - 12, y1 + 29)], fill=arrow)

    # Minimal progress bar.
    bar_x, bar_y, bar_w, bar_h = 1040, 980, 760, 12
@@ -148,10 +148,15 @@ def parse_args():
        default=DURATION_SECONDS,
        help="Animation duration in seconds before final hold. Default: 6",
    )
+    parser.add_argument(
+        "--no-arrow",
+        action="store_true",
+        help="Hide the yellow direction arrow.",
+    )
    return parser.parse_args()


-def generate_video(input_dir, output_path, max_angle=20.0, duration_seconds=6.0):
+def generate_video(input_dir, output_path, max_angle=20.0, duration_seconds=6.0, show_arrow=True):
    output_file = Path(output_path)
    output_file.parent.mkdir(parents=True, exist_ok=True)

@@ -172,17 +177,17 @@ def generate_video(input_dir, output_path, max_angle=20.0, duration_seconds=6.0)
        for index in range(moving_frames):
            t = index / (moving_frames - 1)
            angle = max_angle * smoothstep(t)
-            writer.append_data(np.asarray(make_frame(before_image, angle, max_angle)))
+            writer.append_data(np.asarray(make_frame(before_image, angle, max_angle, show_arrow)))

        for _ in range(hold_frames):
-            writer.append_data(np.asarray(make_frame(before_image, max_angle, max_angle)))
+            writer.append_data(np.asarray(make_frame(before_image, max_angle, max_angle, show_arrow)))

    return output_file.resolve()


 def main():
    args = parse_args()
-    output_file = generate_video(args.input, args.output, args.max_angle, args.duration)
+    output_file = generate_video(args.input, args.output, args.max_angle, args.duration, not args.no_arrow)
    print(output_file)


--- a/web_backend.py
+++ b/web_backend.py
@@ -864,13 +864,14 @@ class Handler(BaseHTTPRequestHandler):
                input_dir = body["inputDir"]
                max_angle = float(body.get("maxAngle", 20))
                duration = float(body.get("durationSeconds", 6))
+                show_arrow = bool(body.get("showArrow", True))

                def worker(job_id):
                    job_root = RESULT_DIR / job_id
                    reset_dir(job_root)
                    output_file = job_root / f"head_extension_{job_id}.mp4"
                    set_job(job_id, message="正在生成 0° 到目标角度的视频。")
-                    output = generate_video(input_dir, output_file, max_angle, duration)
+                    output = generate_video(input_dir, output_file, max_angle, duration, show_arrow)
                    output = Path(output).resolve()
                    return {
                        "video": {