Change image resize to resize and crop

Stability-AI · Dec 22, 2023 · e6a25a8 · e6a25a8
1 parent 32a38ef
commit e6a25a8
Showing 1 changed file with 22 additions and 16 deletions.
diff --git a/nbs/stable_video_v1_REST_API_alpha.ipynb b/nbs/stable_video_v1_REST_API_alpha.ipynb
@@ -47,26 +47,31 @@
         "    image_format = image_mime_type.split(\"/\")[-1].upper()\n",
         "    return image_format\n",
         "\n",
-        "def resize_image_to_bytes(\n",
-        "        image_path : str,\n",
-        "        size : tuple[int,int] = None\n",
-        "):\n",
-        "    # Resize image from file and convert to bytes\n",
-        "    image = Image.open(image_path)\n",
-        "    format = get_image_format(image_path)\n",
-        "    if size is None:\n",
-        "        width, height = get_closest_valid_dims(image)\n",
+        "def resize_and_crop(\n",
+        "            image:Image,\n",
+        "            width:int,\n",
+        "            height:int\n",
+        "    ):\n",
+        "    # Resize the image so one side is the required size, then center crop the other side\n",
+        "    # This is to avoid stretching the image\n",
+        "    w, h = image.size\n",
+        "    if w < h:\n",
+        "        image = image.resize((width, int(h * width / w)))\n",
         "    else:\n",
-        "        width, height = size\n",
-        "    image = image.resize((width, height))\n",
-        "    image_bytes = image_to_bytes(image, format=format)\n",
-        "    return image_bytes\n",
+        "        image = image.resize((int(w * height / h), height))\n",
+        "    # center crop\n",
+        "    left = (image.width - width) / 2\n",
+        "    top = (image.height - height) / 2\n",
+        "    right = (image.width + width) / 2\n",
+        "    bottom = (image.height + height) / 2\n",
+        "    image = image.crop((left, top, right, bottom))\n",
+        "    return image\n",
         "\n",
         "def get_closest_valid_dims(\n",
         "        image : Image\n",
         "):\n",
         "    # Finds the closest aspect ratio to the input image that are valid for SSC\n",
-        "    # Valid dims are 1024x576, 768x768, 1024\n",
+        "    # Valid dimensions are 1024x576, 768x768, 1024\n",
         "    w,h = image.size\n",
         "    aspect_ratio = w/h\n",
         "    portrait_aspect_ratio = 9/16\n",
@@ -78,7 +83,7 @@
         "        width,height = (576,1024) if aspect_ratio < portrait_aspect_ratio_midpoint else (768,768)\n",
         "    else:\n",
         "        # landscape\n",
-        "        width,height = (1024,576) if aspect_ratio < landscape_aspect_ratio_midpoint else (768,768)\n",
+        "        width,height = (1024,576) if aspect_ratio > landscape_aspect_ratio_midpoint else (768,768)\n",
         "\n",
         "    return width, height\n",
         "\n",
@@ -89,7 +94,8 @@
         "    image = Image.open(image_path)\n",
         "    width, height = get_closest_valid_dims(image)\n",
         "    format = get_image_format(image_path)\n",
-        "    image = image.resize((width, height))\n",
+        "    print(f\"Resizing image to {width}x{height}\")\n",
+        "    image = resize_and_crop(image, width, height)\n",
         "    image_bytes = image_to_bytes(image, format=format)\n",
         "    return image_bytes"
       ]