Update app.py

app.py

@@ -7,11 +7,25 @@ import torch
 import numpy as np
 
 def inference(img_1, img_2):
-    # Convert numpy arrays to Tensors
+    # Convert numpy arrays to Tensors and ensure correct shape
     img_1: Tensor = K.image_to_tensor(img_1, keepdim=False).float() / 255.0
-    img_1 = img_1.unsqueeze(0)  # Add batch dimension
     img_2: Tensor = K.image_to_tensor(img_2, keepdim=False).float() / 255.0
-    img_2 = img_2.unsqueeze(0)  # Add batch dimension
+    
+    # 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 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)
+    
+    # Add batch dimension
+    img_1 = img_1.unsqueeze(0)
+    img_2 = img_2.unsqueeze(0)
     
     IS = ImageStitcher(KF.LoFTR(pretrained='outdoor'), estimator='ransac')
     with torch.no_grad():