Coverage for src/rift_console/image_helper.py: 0%

1""" 

2 

3Provides helping functions that are used in image_processing. 

4 

5""" 

6 

7import re 

8import os 

9import datetime 

10from loguru import logger 

11 

12import shared.constants as con 

13from shared.models import CameraAngle 

14 

15 

16def get_angle(image: str) -> CameraAngle: 

17 if "narrow" in image: 

18 return CameraAngle.Narrow 

19 elif "normal" in image: 

20 return CameraAngle.Normal 

21 elif "wide" in image: 

22 return CameraAngle.Wide 

23 logger.warning(f"Unknown camera angle in {image}") 

24 return CameraAngle.Unknown 

25 

26 

27def get_date(image: str) -> str: 

28 # pattern of year-month-dayThour-minute 

29 pattern = r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}" 

30 

31 found_matches = re.findall(pattern, image) 

32 

33 if len(found_matches) == 1: 

34 # logger.debug(found_matches[0]) 

35 match: str = found_matches[0] 

36 return match 

37 else: 

38 logger.warning("None or two dates in {image}") 

39 return datetime.datetime.min.strftime("%Y-%m-%dT%H:%M:%S") 

40 

41 

42def filter_by_date( 

43 images: list[str], start: datetime.datetime, end: datetime.datetime 

44) -> list[str]: 

45 res = [] 

46 date_format = "%Y-%m-%dT%H:%M:%S" 

47 for image in images: 

48 date = datetime.datetime.strptime(get_date(image), date_format).replace( 

49 tzinfo=datetime.timezone.utc 

50 ) 

51 # logger.warning(f"{date} {start} {end}") 

52 if date >= start and date <= end: 

53 res.append(image) 

54 return res 

55 

56 

57def generate_spiral_walk(n: int) -> list[tuple[int, int]]: 

58 """Create an spiraling offset pattern arround a central point, e.g. (0,0), (0,1), (1,0), (1,1), ... 

59 sorted by Manhattan geometry 

60 

61 Args: 

62 n (int): number of offsets to be generated 

63 

64 Returns: 

65 list[tuple[int, int]]: list of offsets 

66 """ 

67 

68 # move right, up, left, down 

69 directions = [(1, 0), (0, 1), (-1, 0), (0, -1)] 

70 # start with going right 

71 direction_index = 0 

72 

73 x, y = 0, 0 

74 offsets = [(x, y)] 

75 

76 # Number of steps we take before changing direction 

77 steps = 1 

78 

79 while len(offsets) < n: 

80 for _ in range(2): 

81 for _ in range(steps): 

82 if len(offsets) < n: 

83 # Move in the current direction 

84 dx, dy = directions[direction_index] 

85 x += dx 

86 y += dy 

87 # Add the new position to the spiral 

88 offsets.append((x, y)) 

89 else: 

90 break 

91 # Change direction clockwise 

92 direction_index = (direction_index + 1) % 4 

93 # After moving two directions, we increase the number of steps 

94 steps += 1 

95 

96 sorted_offset = sorted(offsets, key=lambda x: abs(x[0]) + abs(x[1])) 

97 return sorted_offset 

98 

99 

100def parse_image_name(name: str) -> tuple[int, int, int]: 

101 """Parses an image name in the format generated by Melvonaut and extract the relevant properties for stitching 

102 

103 Args: 

104 name (str): file name of the image 

105 

106 Returns: 

107 tuple[int, int, int]: Used lenssize (and therefore if the image should be scaled to this later) 

108 and approximated x/y coordinates on the stiched image 

109 """ 

110 from shared.models import CameraAngle 

111 

112 # expected format: 'image_5344_wide_2024-12-11T17:31:27.507376_x_19936_y_4879' 

113 # with 8 underscores 

114 if len(name.split("_")) != con.IMAGE_NAME_UNDERSCORE_COUNT: 

115 raise Exception("parse_image_name: filename has wrong format!") 

116 

117 # used CameraAngle is after second underscore 

118 match name.split("_")[con.IMAGE_ANGLE_POSITION]: 

119 case CameraAngle.Narrow: 

120 lens_size = 600 

121 case CameraAngle.Normal: 

122 lens_size = 800 

123 case CameraAngle.Wide: 

124 lens_size = 1000 

125 

126 # find x and y in name 

127 match = re.search(r"_x_(-?\d+)_y_(-?\d+)", name) 

128 

129 if match: 

130 x = int(match.group(1)) 

131 y = int(match.group(2)) 

132 

133 # old images position is not adjusted in melvonaut yet 

134 if con.USE_LEGACY_IMAGE_NAMES: 

135 x -= (int)(lens_size / 2) 

136 y -= (int)(lens_size / 2) 

137 else: 

138 raise Exception("parse_image_name: could not match x/y coordinates!") 

139 

140 return lens_size, x, y 

141 

142 

143# returns all images 

144def find_image_names(directory: str) -> list[str]: 

145 """Traverses the given directory and find + sorts all images in our filename format 

146 

147 Args: 

148 directory (str): path to the folder, needs to include con.IMAGE_PATH 

149 

150 Returns: 

151 list[str]: the name of all images in that folder, sorted by its timestamp from old to now 

152 """ 

153 

154 # find all names 

155 image_names = [] 

156 for filename in os.listdir(directory): 

157 if filename.startswith("image"): 

158 image_names.append(filename) 

159 

160 # helper function used in sorting 

161 def extract_timestamp(s: str) -> datetime.datetime: 

162 timestamp_pattern = r"_(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{6})" 

163 

164 match = re.search(timestamp_pattern, s) 

165 if match: 

166 return datetime.datetime.fromisoformat(match.group(1)) 

167 else: 

168 raise Exception("find_image_names: did not found timestamp in image names") 

169 

170 def extract_pos(s: str) -> int: 

171 pos_pattern = r"_x_(-?\d+)_y_(-?\d+)" 

172 

173 match = re.search(pos_pattern, s) 

174 if match: 

175 x = int(match.group(1)) 

176 y = int(match.group(2)) 

177 return x + y 

178 else: 

179 raise Exception("find_image_names: did not found position in image names") 

180 

181 # sort 

182 if con.SORT_IMAGE_BY_POSITION: 

183 image_names = sorted(image_names, key=extract_pos) 

184 else: 

185 image_names = sorted(image_names, key=extract_timestamp) 

186 return image_names