Update app.py

app.py

@@ -7,23 +7,41 @@ import torch
 import numpy as np
 
 def preprocess_image(img):
+ print(f"Input image type: {type(img)}")
+ print(f"Input image shape: {img.shape if hasattr(img, 'shape') else 'No shape attribute'}")
+ 
  # 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
+ elif isinstance(img, torch.Tensor):
+ img = img.float()
+ if img.max() > 1.0:
+ img = img / 255.0
+ else:
+ raise ValueError(f"Unsupported image type: {type(img)}")
+ 
+ print(f"After conversion to tensor - shape: {img.shape}")
  
  # Ensure 3D tensor (C, H, W)
  if img.ndim == 2:
  img = img.unsqueeze(0)
+ elif img.ndim == 3 and img.shape[0] not in [1, 3]:
+ img = img.permute(2, 0, 1)
+ 
+ print(f"After ensuring 3D - shape: {img.shape}")
  
  # Ensure 3 channel image
  if img.shape[0] == 1:
- img = img.repeat(3, 1, 1)
+ img = img.expand(3, -1, -1)
  elif img.shape[0] > 3:
  img = img[:3] # Take only the first 3 channels if more than 3
  
+ print(f"After ensuring 3 channels - shape: {img.shape}")
+ 
  # Add batch dimension
  img = img.unsqueeze(0)
  
+ print(f"Final tensor shape: {img.shape}")
  return img
 
 def inference(img_1, img_2):