Spaces:
Runtime error
Runtime error
File size: 4,057 Bytes
f4998a2 ac7cda5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
try:
import filetype
def is_image(file_path):
return filetype.is_image(file_path)
def is_video(file_path):
return filetype.is_video(file_path)
except ImportError:
# Fallback implementation if filetype is not available
import os
def is_image(file_path):
"""Check if file is an image based on extension"""
if not os.path.exists(file_path):
return False
ext = os.path.splitext(file_path)[1].lower()
return ext in ['.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff', '.webp']
def is_video(file_path):
"""Check if file is a video based on extension"""
if not os.path.exists(file_path):
return False
ext = os.path.splitext(file_path)[1].lower()
return ext in ['.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv', '.webm', '.m4v']
import imageio
import cv2
def check_resize(h, w, max_dim=1920, division=2):
rsz_flag = False
# ajust the size of the image according to the maximum dimension
if max_dim > 0 and max(h, w) > max_dim:
rsz_flag = True
if h > w:
new_h = max_dim
new_w = int(round(w * max_dim / h))
else:
new_w = max_dim
new_h = int(round(h * max_dim / w))
else:
new_h = h
new_w = w
# ensure that the image dimensions are multiples of n
if new_h % division != 0:
new_h = new_h - (new_h % division)
rsz_flag = True
if new_w % division != 0:
new_w = new_w - (new_w % division)
rsz_flag = True
return new_h, new_w, rsz_flag
def load_image(image_path, max_dim=-1):
img = cv2.imread(image_path, cv2.IMREAD_COLOR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
h, w = img.shape[:2]
new_h, new_w, rsz_flag = check_resize(h, w, max_dim)
if rsz_flag:
img = cv2.resize(img, (new_w, new_h))
return img
def load_video(video_path, n_frames=-1, max_dim=-1):
reader = imageio.get_reader(video_path, "ffmpeg")
new_h, new_w, rsz_flag = None, None, None
ret = []
for idx, frame_rgb in enumerate(reader):
if n_frames > 0 and idx >= n_frames:
break
if rsz_flag is None:
h, w = frame_rgb.shape[:2]
new_h, new_w, rsz_flag = check_resize(h, w, max_dim)
if rsz_flag:
frame_rgb = cv2.resize(frame_rgb, (new_w, new_h))
ret.append(frame_rgb)
reader.close()
return ret
def load_source_frames(source_path, max_dim=-1, n_frames=-1):
if is_image(source_path):
rgb = load_image(source_path, max_dim)
rgb_list = [rgb]
is_image_flag = True
elif is_video(source_path):
rgb_list = load_video(source_path, n_frames, max_dim)
is_image_flag = False
else:
raise ValueError(f"Unsupported source type: {source_path}")
return rgb_list, is_image_flag
def _mirror_index(index, size):
turn = index // size
res = index % size
if turn % 2 == 0:
return res
else:
return size - res - 1
class LoopLoader:
def __init__(self, item_list, max_iter_num=-1, mirror_loop=True):
self.item_list = item_list
self.idx = 0
self.item_num = len(self.item_list)
self.max_iter_num = max_iter_num if max_iter_num > 0 else self.item_num
self.mirror_loop = mirror_loop
def __len__(self):
return self.max_iter_num
def __iter__(self):
return self
def __next__(self):
if self.idx >= self.max_iter_num:
raise StopIteration
if self.mirror_loop:
idx = _mirror_index(self.idx, self.item_num)
else:
idx = self.idx % self.item_num
item = self.item_list[idx]
self.idx += 1
return item
def __call__(self):
return self.__iter__()
def reset(self, max_iter_num=-1):
self.frame_idx = 0
self.max_iter_num = max_iter_num if max_iter_num > 0 else self.item_num
|