Skip to content

Commit

Permalink
added upright option from exif
Browse files Browse the repository at this point in the history
  • Loading branch information
lcmrl committed Sep 9, 2024
1 parent 19ebf53 commit 4e1dc1f
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 45 deletions.
2 changes: 1 addition & 1 deletion src/deep_image_matching/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -243,7 +243,7 @@
"matching_lowres",
"covisibility",
],
"upright_strategy": ["custom", "2-clusters"]
"upright_strategy": ["custom", "2clusters", "exif"]
}


Expand Down
107 changes: 63 additions & 44 deletions src/deep_image_matching/image_matching.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
from pathlib import Path
from pprint import pprint
from PIL import Image
import PIL.ExifTags
import exifread

import cv2
import h5py
Expand Down Expand Up @@ -39,6 +41,7 @@
from functools import partial

from .utils import ImageList, get_pairs_from_file
from .config import Config

logger = logging.getLogger("dim")
timer = Timer(logger=logger)
Expand Down Expand Up @@ -173,7 +176,7 @@ def upright(
# image1 = cv2.imread(str(path_to_upright_dir / img1))
# rotated_image1 = cv2.rotate(image1, cv2_rot_params[index_of_max])
# rotated_image1 = cv2.cvtColor(rotated_image1, cv2.COLOR_BGR2GRAY)
p = image1.rotate(cv2_rot_params[index_of_max], expand=False)
p = image1.rotate(cv2_rot_params[index_of_max], expand=True)
p.save(str(path_to_upright_dir / img1))
processed.append(img1)
break
Expand Down Expand Up @@ -421,7 +424,8 @@ def rotate_upright_images(

logger.info(f"Copying images to {path_to_upright_dir}")
for img in images:
shutil.copy(self.image_dir / img, path_to_upright_dir / img)
shutil.copy2(self.image_dir / img, path_to_upright_dir / img)

logger.info(f"{len(images)} images copied")

rotations = [0, 90, 180, 270]
Expand All @@ -441,54 +445,38 @@ def rotate_upright_images(
self.rotated_images = []

if strategy == "2clusters":

logger.info(f"Initializing Superpoint + LIghtGlue..")
SPextractor = SuperPointExtractor(
config={
"general": {},
"extractor": {
"keypoint_threshold": 0.005,
"max_keypoints": 1024,
},
}
)
LGmatcher = LightGlueMatcher(
config={
"general": {},
"matcher": {
"depth_confidence": 0.95, # early stopping, disable with -1
"width_confidence": 0.99, # point pruning, disable with -1
"filter_threshold": 0.1, # match threshold
},
},
)
SPextractor = SuperPointExtractor(self.config)
LGmatcher = LightGlueMatcher(self.config)

#SPextractor = SuperPointExtractor(
# config={
# "general": {},
# "extractor": {
# "keypoint_threshold": 0.005,
# "max_keypoints": 1024,
# },
# }
#)
#LGmatcher = LightGlueMatcher(
# config={
# "general": {},
# "matcher": {
# "depth_confidence": 0.95, # early stopping, disable with -1
# "width_confidence": 0.99, # point pruning, disable with -1
# "filter_threshold": 0.1, # match threshold
# },
# },
#)

cluster0 = []
cluster1 = os.listdir(path_to_upright_dir)

SPextractor = SuperPointExtractor(
config={
"general": {},
"extractor": {
"keypoint_threshold": 0.005,
"max_keypoints": 1024,
},
}
)
LGmatcher = LightGlueMatcher(
config={
"general": {},
"matcher": {
"depth_confidence": 0.95, # early stopping, disable with -1
"width_confidence": 0.99, # point pruning, disable with -1
"filter_threshold": 0.1, # match threshold
},
},
)

# Random init
random_first_img = random.randint(0, len(cluster1))
# Choose first image
# random_first_img = 1
# random_first_img = 0

cluster0.append(cluster1[random_first_img])
cluster1.pop(random_first_img)
Expand Down Expand Up @@ -571,6 +559,36 @@ def rotate_upright_images(
except:
pass

if strategy == "exif":
orientation_map = {
'Horizontal (normal)': 0,
'Rotated 180': 180,
'Rotated 90 CW': 90,
'Rotated 90 CCW': 270
}

for img in os.listdir(self.image_dir):
image_path = path_to_upright_dir / img
image = cv2.imread(str(self.image_dir / img))

with open(str(self.image_dir / img), 'rb') as image_file:
tags = exifread.process_file(image_file)
orientation_tag = 'Image Orientation'

if orientation_tag in tags:
orientation_description = str(tags[orientation_tag])
orientation_degrees = orientation_map.get(orientation_description, None)
print(orientation_degrees)
if orientation_degrees is not None:
if orientation_degrees == 180:
image = cv2.rotate(image, cv2.ROTATE_180)
elif orientation_degrees == 90:
image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
elif orientation_degrees == 270:
image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
cv2.imwrite(str(image_path), image)


out_file = self.pair_file.parent / f"{self.pair_file.stem}_rot.txt"
with open(out_file, "w") as txt_file:
for element in self.rotated_images:
Expand Down Expand Up @@ -682,6 +700,7 @@ def rotate_back_features(self, feature_path: Path) -> None:
"""
# images = self.image_list.img_names
for img, theta in tqdm(self.rotated_images):
print('img, theta', img, theta)
features = get_features(feature_path, img)
keypoints = features["keypoints"]
rotated_keypoints = np.empty(keypoints.shape)
Expand All @@ -702,14 +721,14 @@ def rotate_back_features(self, feature_path: Path) -> None:
x_rot = W - x
rotated_keypoints[r, 0], rotated_keypoints[r, 1] = x_rot, y_rot

if theta == 270:
if theta == 90:
for r in range(keypoints.shape[0]):
x, y = keypoints[r, 0], keypoints[r, 1]
y_rot = W - x
x_rot = y
rotated_keypoints[r, 0], rotated_keypoints[r, 1] = x_rot, y_rot

if theta == 90:
if theta == 270:
for r in range(keypoints.shape[0]):
x, y = keypoints[r, 0], keypoints[r, 1]
y_rot = x
Expand Down

0 comments on commit 4e1dc1f

Please sign in to comment.