Update app.py

app.py

@@ -6,26 +6,30 @@ import kornia.feature as KF
 import torch
 import numpy as np
 
-def inference(img_1, img_2):
-    # Convert numpy arrays to Tensors and ensure correct shape
-    img_1: Tensor = K.image_to_tensor(img_1, keepdim=False).float() / 255.0
-    img_2: Tensor = K.image_to_tensor(img_2, keepdim=False).float() / 255.0
+def preprocess_image(img):
+    # Convert numpy array to Tensor and ensure correct shape
+    if isinstance(img, np.ndarray):
+        img = K.image_to_tensor(img, keepdim=False).float() / 255.0
     
-    # Ensure 3D tensors (C, H, W)
-    if img_1.ndim == 2:
-        img_1 = img_1.unsqueeze(0)
-    if img_2.ndim == 2:
-        img_2 = img_2.unsqueeze(0)
+    # Ensure 3D tensor (C, H, W)
+    if img.ndim == 2:
+        img = img.unsqueeze(0)
     
-    # Ensure 3 channel images
-    if img_1.shape[0] == 1:
-        img_1 = img_1.repeat(3, 1, 1)
-    if img_2.shape[0] == 1:
-        img_2 = img_2.repeat(3, 1, 1)
+    # Ensure 3 channel image
+    if img.shape[0] == 1:
+        img = img.repeat(3, 1, 1)
+    elif img.shape[0] > 3:
+        img = img[:3]  # Take only the first 3 channels if more than 3
     
     # Add batch dimension
-    img_1 = img_1.unsqueeze(0)
-    img_2 = img_2.unsqueeze(0)
+    img = img.unsqueeze(0)
+    
+    return img
+
+def inference(img_1, img_2):
+    # Preprocess images
+    img_1 = preprocess_image(img_1)
+    img_2 = preprocess_image(img_2)
     
     IS = ImageStitcher(KF.LoFTR(pretrained='outdoor'), estimator='ransac')
     with torch.no_grad():