Add support for WebP images

demo/configs/webpack/common.js

@@ -1,9 +1,3 @@
-// Copyright (c) Meta Platforms, Inc. and affiliates.
-// All rights reserved.
-
-// This source code is licensed under the license found in the
-// LICENSE file in the root directory of this source tree.
-
 const { resolve } = require("path");
 const HtmlWebpackPlugin = require("html-webpack-plugin");
 const FriendlyErrorsWebpackPlugin = require("friendly-errors-webpack-plugin");
@@ -42,10 +36,24 @@ module.exports = {
         use: ["style-loader", "css-loader", "postcss-loader"],
       },
       {
-        test: /\.(jpe?g|png|gif|svg)$/i,
+        test: /\.(jpe?g|png|gif|svg|webp)$/i,
         use: [
           "file-loader?hash=sha512&digest=hex&name=img/[contenthash].[ext]",
-          "image-webpack-loader?bypassOnDebug&optipng.optimizationLevel=7&gifsicle.interlaced=false",
+          {
+            loader: "image-webpack-loader",
+            options: {
+              bypassOnDebug: true,
+              optipng: {
+                optimizationLevel: 7,
+              },
+              gifsicle: {
+                interlaced: false,
+              },
+              webp: {
+                quality: 75,
+              },
+            },
+          },
         ],
       },
       {

demo/src/App.tsx

@@ -1,9 +1,3 @@
-// Copyright (c) Meta Platforms, Inc. and affiliates.
-// All rights reserved.
-
-// This source code is licensed under the license found in the
-// LICENSE file in the root directory of this source tree.
-
 import { InferenceSession, Tensor } from "onnxruntime-web";
 import React, { useContext, useEffect, useState } from "react";
 import "./assets/scss/App.scss";
@@ -76,6 +70,11 @@ const App = () => {
         img.height = height; 
         setImage(img);
       };
+      img.onerror = () => {
+        if (url.href.endsWith(".webp")) {
+          console.log("Failed to load WebP image. Please ensure the image is in the correct format.");
+        }
+      };
     } catch (error) {
       console.log(error);
     }

scripts/amg.py

@@ -212,11 +212,14 @@ def main(args: argparse.Namespace) -> None:
 
     for t in targets:
         print(f"Processing '{t}'...")
-        image = cv2.imread(t)
+        image = cv2.imread(t, cv2.IMREAD_UNCHANGED)
         if image is None:
             print(f"Could not load '{t}' as an image, skipping...")
             continue
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        if t.lower().endswith('.webp'):
+            image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA)
+        else:
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
         masks = generator.generate(image)
 

scripts/export_onnx_model.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
+import cv2  # type: ignore
 
 from segment_anything import sam_model_registry
 from segment_anything.utils.onnx import SamOnnxModel
@@ -173,6 +174,13 @@ def to_numpy(tensor):
 
 if __name__ == "__main__":
     args = parser.parse_args()
+    image = cv2.imread(args.input, cv2.IMREAD_UNCHANGED)
+    if image is None:
+        print(f"Could not load '{args.input}' as an image, skipping...")
+    if args.input.lower().endswith('.webp'):
+        image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA)
+    else:
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     run_export(
         model_type=args.model_type,
         checkpoint=args.checkpoint,