Update inference.py

inference.py

@@ -268,7 +268,7 @@ def main():
 
     if not args.audio.endswith('.wav'):
         print('Extracting raw audio...')
-        temp_wav = 'temp/temp.wav'
+        temp_wav = os.path.join(os.path.dirname(args.outfile), 'temp.wav')
         command = f'ffmpeg -y -i "{args.audio}" -strict -2 "{temp_wav}"'
         subprocess.call(command, shell=True)
         args.audio = temp_wav
@@ -307,11 +307,13 @@ def main():
     if args.save_as_video:
         frame_sample = next(reader)
         frame_h, frame_w = frame_sample.shape[:2]
-        out = cv2.VideoWriter('temp/result.avi', 
+        # Определяем путь для result.avi в той же директории, что и outfile
+        result_avi = os.path.join(os.path.dirname(args.outfile), "result.avi")
+        out = cv2.VideoWriter(result_avi, 
                               cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h))
         if args.save_frames:
-            gt_out = cv2.VideoWriter("temp/gt.avi", cv2.VideoWriter_fourcc(*'DIVX'), fps, (384, 384))
-            pred_out = cv2.VideoWriter("temp/pred.avi", cv2.VideoWriter_fourcc(*'DIVX'), fps, (96, 96))
+            gt_out = cv2.VideoWriter(os.path.join(os.path.dirname(args.outfile), "gt.avi"), cv2.VideoWriter_fourcc(*'DIVX'), fps, (384, 384))
+            pred_out = cv2.VideoWriter(os.path.join(os.path.dirname(args.outfile), "pred.avi"), cv2.VideoWriter_fourcc(*'DIVX'), fps, (96, 96))
     else:
         out = None
         gt_out = None
@@ -327,7 +329,7 @@ def main():
             pred = model(mel_batch, img_batch)
 
         pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.0
-        
+
         for p, f, c in zip(pred, frames, coords):
             y1, y2, x1, x2 = c
 
@@ -364,16 +366,17 @@ def main():
     if out:
         out.release()
 
-    # Объединение аудио и видео
-    final_command = f'ffmpeg -y -i "{args.audio}" -i "temp/result.avi" -strict -2 -q:v 1 "{args.outfile}"'
-    subprocess.call(final_command, shell=(platform.system() != 'Windows'))
-
+    # Определение пути к result.avi
+    if args.save_as_video:
+        final_command = f'ffmpeg -y -i "{args.audio}" -i "{result_avi}" -strict -2 -q:v 1 "{args.outfile}"'
+        subprocess.call(final_command, shell=(platform.system() != 'Windows'))
+    
     if args.save_frames and args.save_as_video:
         gt_out.release()
         pred_out.release()
 
-        gt_video_cmd = f'ffmpeg -y -i "temp/gt.avi" -i "{args.audio}" -strict -2 -q:v 1 "{args.gt_path}"'
-        pred_video_cmd = f'ffmpeg -y -i "temp/pred.avi" -i "{args.audio}" -strict -2 -q:v 1 "{args.pred_path}"'
+        gt_video_cmd = f'ffmpeg -y -i "{os.path.join(os.path.dirname(args.outfile), "gt.avi")}" -i "{args.audio}" -strict -2 -q:v 1 "{args.gt_path}"'
+        pred_video_cmd = f'ffmpeg -y -i "{os.path.join(os.path.dirname(args.outfile), "pred.avi")}" -i "{args.audio}" -strict -2 -q:v 1 "{args.pred_path}"'
 
         subprocess.call(gt_video_cmd, shell=(platform.system() != 'Windows'))
         subprocess.call(pred_video_cmd, shell=(platform.system() != 'Windows'))