Update my_model/state_manager.py

my_model/state_manager.py

@@ -256,47 +256,46 @@ class StateManager:
         """
         return st.session_state['images_data']
 
-
-
-    def resize_image_aspect_ratio(self,image_input, base_width=None, base_height=None):
-        """
-        Resize an image while maintaining its aspect ratio.
-    
-        Args:
-        image (PIL.Image.Image): The image to resize.
-        base_width (int, optional): The target width, with the height adjusted to maintain aspect ratio.
-        base_height (int, optional): The target height, with the width adjusted to maintain aspect ratio.
-    
-        Returns:
-        PIL.Image.Image: The resized image.
-        """
-
-        image = copy.deepcopy(image_input)
-        if isinstance(image, str):
+def resize_image(image_input, new_width=None, new_height=None):
+    """
+    Resize an image. If only new_width is provided, the height is adjusted to maintain aspect ratio.
+    If both new_width and new_height are provided, the image is resized to those dimensions.
+
+    Args:
+    image (PIL.Image.Image): The image to resize.
+    new_width (int, optional): The target width of the image.
+    new_height (int, optional): The target height of the image.
+
+    Returns:
+    PIL.Image.Image: The resized image.
+    """
+
+    img = copy.deepcopy(image_input)
+        if isinstance(img, str):
             # Open the image from a file path
-            image = Image.open(image)
-        elif isinstance(image, Image.Image):
+            image = Image.open(img)
+        elif isinstance(img, Image.Image):
             # Use the image directly if it's already a PIL Image object
-            image = image_input
+            image = img
         else:
             raise ValueError("image_input must be a file path or a PIL Image object")
-
-        # Get original dimensions
+            
+    if new_width is not None and new_height is None:
+        # Calculate new height to maintain aspect ratio
         original_width, original_height = image.size
-    
-        # Calculate new dimensions
-        if base_width is not None:
-            ratio = (base_width / float(original_width))
-            new_height = int((float(original_height) * float(ratio)))
-            resized_image = image.resize((base_width, new_height))
-        elif base_height is not None:
-            ratio = (base_height / float(original_height))
-            new_width = int((float(original_width) * float(ratio)))
-            resized_image = image.resize((new_width, base_height))
-        else:
-            raise ValueError("Either base_width or base_height must be provided")
-    
-        return resized_image
+        ratio = new_width / original_width
+        new_height = int(original_height * ratio)
+    elif new_width is None and new_height is not None:
+        # Calculate new width to maintain aspect ratio
+        original_width, original_height = image.size
+        ratio = new_height / original_height
+        new_width = int(original_width * ratio)
+    elif new_width is None and new_height is None:
+        raise ValueError("At least one of new_width or new_height must be provided")
+
+    # Resize the image
+    resized_image = image.resize((new_width, new_height))
+    return resized_image