6
0

file_ops.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141
  1. import cv2
  2. import os
  3. from PIL import Image
  4. from pycs.database.File import File
  5. def crop_file(file_id: int, x: float, y: float, w: float, h: float) -> str:
  6. """
  7. gets a file for the given file_id, crops the according image to the
  8. bounding box and saves the crops in the temp folder of the project.
  9. :param file: file object
  10. :param x: relative x-coordinate of the top left corner
  11. :param y: relative y-coordinate of the top left corner
  12. :param w: relative width of the bounding box
  13. :param h: relative height of the bounding box
  14. :return: directory and file name of the cropped patch
  15. """
  16. file = File.query.get(file_id)
  17. file_path = file.absolute_path
  18. root_folder = file.project.root_folder
  19. image = Image.open(file_path)
  20. width, height = image.size
  21. crop_width = int(width * w)
  22. crop_height = int(height * h)
  23. x0 = int(width * x)
  24. y0 = int(height * y)
  25. x1 = x0 + crop_width
  26. y1 = y0 + crop_height
  27. target_path = os.path.join(
  28. os.getcwd(),
  29. root_folder,
  30. 'temp',
  31. f'{file.uuid}_{x0}-{y0}_{x1}-{y1}.jpg',
  32. )
  33. if not os.path.exists(target_path):
  34. crop = image.crop((x0, y0, x0+crop_width, y0+crop_height))
  35. crop.save(target_path, quality=80)
  36. return os.path.split(target_path)
  37. def resize_file(file_id: int, max_width: int, max_height: int) -> tuple[str, str]:
  38. """
  39. If file type equals video this function extracts a thumbnail first. It calls resize_image
  40. to resize and returns the resized files directory and name.
  41. :param file: file object
  42. :param max_width: maximum image or thumbnail width
  43. :param max_height: maximum image or thumbnail height
  44. :return: resized file directory, resized file name
  45. """
  46. file = File.query.get(file_id)
  47. project = file.project
  48. abs_file_path = file.absolute_path
  49. # extract video thumbnail
  50. if file.type == 'video':
  51. abs_target_path = os.path.join(os.getcwd(), project.root_folder, 'temp', f'{file.uuid}.jpg')
  52. create_thumbnail(abs_file_path, abs_target_path)
  53. abs_file_path = abs_target_path
  54. # resize image file
  55. abs_target_path = os.path.join(os.getcwd(), project.root_folder,
  56. 'temp', f'{file.uuid}_{max_width}_{max_height}.jpg')
  57. result = resize_image(abs_file_path, abs_target_path, max_width, max_height)
  58. # return path
  59. if result is not None:
  60. return os.path.split(abs_target_path)
  61. return os.path.split(abs_file_path)
  62. def resize_image(file_path: str, target_path: str, max_width: int, max_height: int):
  63. """
  64. resize an image so width < max_width and height < max_height
  65. :param file_path: path to source file
  66. :param target_path: path to target file
  67. :param max_width: maximum image width
  68. :param max_height: maximum image height
  69. :return:
  70. """
  71. # return if file exists
  72. if os.path.exists(target_path):
  73. return True
  74. # load full size image
  75. image = Image.open(file_path)
  76. img_width, img_height = image.size
  77. # abort if file is smaller than desired
  78. if img_width < max_width and img_height < max_height:
  79. return None
  80. # calculate target size
  81. target_width = int(max_width)
  82. target_height = int(max_width * img_height / img_width)
  83. if target_height > max_height:
  84. target_height = int(max_height)
  85. target_width = int(max_height * img_width / img_height)
  86. # resize image
  87. resized_image = image.resize((target_width, target_height))
  88. # save to file
  89. resized_image.save(target_path, quality=80)
  90. return True
  91. def create_thumbnail(file_path: str, target_path: str):
  92. """
  93. extract a thumbnail from a video
  94. :param file_path: path to source file
  95. :param target_path: path to target file
  96. :return:
  97. """
  98. # return if file exists
  99. if os.path.exists(target_path):
  100. return
  101. # load video
  102. video = cv2.VideoCapture(file_path)
  103. # create thumbnail
  104. _, image = video.read()
  105. cv2.imwrite(target_path, image)
  106. # close video file
  107. video.release()