цикл Python не вызывает функцию во второй раз

Question

Ниже код дает оценку позы «лица» в видео.Я изменил код, чтобы взять папку / каталог в качестве входных данных и ожидать, что он будет обрабатывать все видео в каталоге.Используя приведенный ниже код, я ожидаю, что все видео в папке будут обработаны, но цикл for обработает только одно видео, а не другие, ниже представлен цикл, и он вызовет parse_video только один раз.

    if args.videoDirPath is not None:
    for videoName in os.listdir(folderName):
        print(videoName)
        video = cv2.VideoCapture(videoName)
        parse_video(video)

Папка (videoFolder) имеет следующие видео:

amir.mp4
arnab-srk.mp4
kanihya.mp4
simma.mp4
salman.mp4

вывод

opt/anaconda3/lib/python3.7/site- 
packages/torchvision/transforms/transforms.py:207: UserWarning: The use of 
the transforms.Scale transform is deprecated, please use transforms.Resize 
instead.
warnings.warn("The use of the transforms.Scale transform is deprecated, " +

simma.mp4
frameNumber : 1
amir.mp4
creating...output/frame1.jpg
creating...output/frame2.jpg
creating...output/frame3.jpg
creating...output/frame4.jpg
creating...output/frame5.jpg
frameNumber : 6
arnab-srk.mp4
frameNumber : 6
kanihya.mp4
frameNumber : 6
salman.mp4
frameNumber : 6

Выходная папка: содержит следующие видео и текстовый файл:

     output-out-1.avi
     output-out-6.avi
     output-out.txt  # blank

Я запускаю программу, используя следующие параметры

!python code/test_on_video_dlib.py --snapshot hopenet_alpha1.pkl --face_model mmod_human_face_detector.dat --directoryPath videoFolder --output_string out --n_frames 20 --fps 200enter code here  

Код для 'test_on_video_dlib.py'

  import sys, os, argparse
  import numpy as np
  import cv2
  import matplotlib.pyplot as plt
  import torch
  import torch.nn as nn
  from torch.autograd import Variable
  from torch.utils.data import DataLoader
  from torchvision import transforms
  import torch.backends.cudnn as cudnn
  import torchvision
  import torch.nn.functional as F
  from PIL import Image
  import datasets, hopenet, utils
  from skimage import io
  import dlib
  import face_alignment
  import numpy as np
  from mpl_toolkits.mplot3d import Axes3D
  import matplotlib.pyplot as plt
  from skimage import io

  def parse_video(video,nr):       

      # New cv2
      width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))   # float
      height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) # float

      # Define the codec and create VideoWriter object
      fourcc = cv2.VideoWriter_fourcc(*'MJPG')

      out = cv2.VideoWriter('output/video/output-{}-{}.avi'.format(args.output_string, nr), fourcc, 
                  args.fps, (width, height))
      #frame_num = 1

      frame_num = nr # add nr here also

      while frame_num <= args.n_frames:
          #print frame_num
          ret,frame = video.read()
          if ret == False:
              break            
          #writing frames
          name = 'output/frame' + str(frame_num) + '.jpg'
          print("creating..." +name)
          cv2.imwrite(name,frame)        
          cv2_frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
          # Dlib detect
          dets = cnn_face_detector(cv2_frame, 1)

          for idx, det in enumerate(dets):
            # Get x_min, y_min, x_max, y_max, conf
            x_min = det.rect.left()
            y_min = det.rect.top()
            x_max = det.rect.right()
            y_max = det.rect.bottom()
            conf = det.confidence

            if conf > 1.0:
                bbox_width = abs(x_max - x_min)
                bbox_height = abs(y_max - y_min)
                x_min -= 2 * bbox_width / 4
                x_max += 2 * bbox_width / 4
                y_min -= 3 * bbox_height / 4
                y_max += bbox_height / 4
                x_min = max(x_min, 0); y_min = max(y_min, 0)
                x_max = min(frame.shape[1], x_max); y_max = min(frame.shape[0], y_max)
                # Crop image
                img = cv2_frame[int(y_min):int(y_max),int(x_min):int(x_max)]
                img = Image.fromarray(img)

                # Transform
                img = transformations(img)
                img_shape = img.size()
                img = img.view(1, img_shape[0], img_shape[1], img_shape[2])
                img = Variable(img).cuda(gpu)

                yaw, pitch, roll = model(img)

                yaw_predicted = F.softmax(yaw,dim=1)
                pitch_predicted = F.softmax(pitch,dim=1)
                roll_predicted = F.softmax(roll,dim=1)
                # Get continuous predictions in degrees.
                yaw_predicted = torch.sum(yaw_predicted.data[0] * idx_tensor) * 3 - 99
                pitch_predicted = torch.sum(pitch_predicted.data[0] * idx_tensor) * 3 - 99
                roll_predicted = torch.sum(roll_predicted.data[0] * idx_tensor) * 3 - 99
                txt_out.write(('output/frame' + str(frame_num) + '.jpg') + ' %f %f %f\n' % (yaw_predicted, pitch_predicted,       roll_predicted))
                # utils.plot_pose_cube(frame, yaw_predicted, pitch_predicted, roll_predicted, (x_min + x_max) / 2, (y_min + y_max) / 2, size = bbox_width)
                utils.draw_axis(frame, yaw_predicted, pitch_predicted, roll_predicted, tdx = (x_min + x_max) / 2, tdy= (y_min + y_max) / 2, size = bbox_height/2)
                # Plot expanded bounding box
                # cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (0,255,0), 1)

       out.write(frame)
       frame_num += 1

  out.release()
  video.release()
  return frame_num

  def parse_args():
      """Parse input arguments."""
      parser = argparse.ArgumentParser(description='Head pose estimation using the Hopenet network.')
      parser.add_argument('--gpu', dest='gpu_id', help='GPU device id to use [0]',
        default=0, type=int)
      parser.add_argument('--snapshot', dest='snapshot', help='Path of model snapshot.',
      default='', type=str)
      parser.add_argument('--face_model', dest='face_model', help='Path of DLIB face detection model.',
      default='', type=str)
      parser.add_argument('--video', dest='video_path', help='Path of video')
      #code to pass video folder name
      parser.add_argument('--directoryPath',dest='videoDirPath' ,help="directory path containing all videos")
      parser.add_argument('--output_string', dest='output_string', help='String appended to output file')
      parser.add_argument('--n_frames', dest='n_frames', help='Number of frames', type=int)
      parser.add_argument('--fps', dest='fps', help='Frames per second of source video', type=float, default=30.)
      args = parser.parse_args()
      return args

  if __name__ == '__main__':
      args = parse_args()
      cudnn.enabled = True

      batch_size = 1
      gpu = args.gpu_id
      snapshot_path = args.snapshot
      out_dir = 'output/video'
      video_path = args.video_path
      #folder path code
      folderName = args.videoDirPath

      if not os.path.exists(out_dir):
          os.makedirs(out_dir)

      # ResNet50 structure
      model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 66)

      # Dlib face detection model
      cnn_face_detector = dlib.cnn_face_detection_model_v1(args.face_model)

      #print 'Loading snapshot.'
      # Load snapshot
      saved_state_dict = torch.load(snapshot_path)
      model.load_state_dict(saved_state_dict)

      #print 'Loading data.'

      transformations = transforms.Compose([transforms.Scale(224),
      transforms.CenterCrop(224), transforms.ToTensor(),
      transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

      model.cuda(gpu)

      #print 'Ready to test network.'

      # Test the Model
       model.eval()  # Change model to 'eval' mode (BN uses moving mean/var).
      total = 0

      idx_tensor = [idx for idx in range(66)]
      idx_tensor = torch.FloatTensor(idx_tensor).cuda(gpu)

      if args.video_path is not None:
              video = cv2.VideoCapture(video_path)
              parse_video(video)
      # THIS IS THE LOOP I AM REFERRING IN QUESTION   
      nr=1 
      if args.videoDirPath is not None:
         for videoName in os.listdir(folderName):
             print(videoName)
             video = cv2.VideoCapture(videoName)
             nr = parse_video(video ,nr)    

Ожидаемый результат:

Я хочу, чтобы каждое видео в videoFolder было обработано, а его кадр должен быть создан в выходной папке.

Answer 1

Что касается меня, вы должны использовать правильный путь к файлу - folderName/videoName

for videoName in os.listdir(folderName):

    videoName = os.path.join(folderName, videoName)  

    print(videoName)
    video = cv2.VideoCapture(videoName)