import random as rd
import cv2 as cv
import numpy as np


# 保存視頻
class RecordMovie(object):

  def __init__(self, img_width, img_height):
    self.video_writer = None # 視頻對象
    self.is_end = False # 結(jié)束保存視頻
    self.img_width = img_width # 寬度
    self.img_height = img_height # 高度

  # 創(chuàng)建 視頻寫入對象
  def start(self, file_name, freq):
    # 創(chuàng)建視頻格式
    four_cc = cv.VideoWriter_fourcc(*'mp4v')
    img_size = (self.img_width, self.img_height) # 視頻尺寸

    # 創(chuàng)建視頻寫入對象
    self.video_writer = cv.VideoWriter()
    self.video_writer.open(file_name, four_cc, freq, img_size, True)

  # 寫入圖片幀
  def record(self, img):
    if self.is_end is False:
      self.video_writer.write(img)

  # 完成視頻 釋放資源
  def end(self):
    self.is_end = True
    self.video_writer.release()


def move_image(img_src):
  img_height, img_width = img_src.shape[:2]

  # 隨機(jī) xy平移方向與大小設(shè)置
  x_size = rd.randint(-3, 3)
  y_size = rd.randint(-3, 3)

  # 自定義轉(zhuǎn)換矩陣
  transform_matrix = np.float32([[1, 0, x_size], [0, 1, y_size]])

  # 執(zhí)行平移
  return cv.warpAffine(img_src, transform_matrix, (img_width, img_height))


def main():
  # 1.讀取圖片
  img_org = cv.imread("img.png", cv.IMREAD_GRAYSCALE)

  # 2.顯示圖片
  cv.imshow("org", img_org)
  cv.namedWindow("shift")

  # 3.視頻文件生成
  height, width = img_org.shape[:2]
  print(height, width)
  rm = RecordMovie(width, height)

  # 設(shè)置視頻文件名稱 頻率
  rm.start("test.mp4", 20)

  # 4.圖片寫入視頻
  for i in range(300):
    # 圖片微調(diào)調(diào)整
    img_move = move_image(img_org)
    img_move = cv.cvtColor(img_move, cv.COLOR_GRAY2RGB)

    rm.record(img_move)
    cv.imshow("shift", img_move)
    key = cv.waitKey(10)
    if key == 27: # esc 按鍵
      break

  # 5.關(guān)閉視頻文件
  rm.end()


if __name__ == '__main__':
  main()