diff --git a/requirements.txt b/requirements.txt index f6f70d2..9029ff3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,3 +9,4 @@ tqdm==4.51.0 requests==2.24.0 scipy==1.5.4 pymatting==1.1.1 +filetype=1.0.7 diff --git a/src/rembg/cmd/cli.py b/src/rembg/cmd/cli.py index ce0229a..b13c2e6 100644 --- a/src/rembg/cmd/cli.py +++ b/src/rembg/cmd/cli.py @@ -2,6 +2,7 @@ import argparse import glob import imghdr import os +import filetype from distutils.util import strtobool from ..bg import remove @@ -90,7 +91,10 @@ def main(): full_paths += glob.glob(path + "/*") for fi in files: - if imghdr.what(fi) is None: + fi_type = filetype.guess(fi) + if fi_type is None: + continue + elif if_type.mime.find('image') < 0: continue with open(fi, "rb") as input: