AI/rembg
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg synced 2025-08-16 02:25:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

[bug fix] change image file detector

Browse Source
This commit is contained in:
Jaric Kuo 2020-11-13 14:53:23 +08:00
parent 4eb79878dc
commit 8f360f1a01
2 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
requirements.txt
View File

@ -9,3 +9,4 @@ tqdm==4.51.0
requests==2.24.0 requests==2.24.0
scipy==1.5.4 scipy==1.5.4
pymatting==1.1.1 pymatting==1.1.1
filetype=1.0.7

6
src/rembg/cmd/cli.py
View File

@ -2,6 +2,7 @@ import argparse
import glob import glob
import imghdr import imghdr
import os import os
import filetype
from distutils.util import strtobool from distutils.util import strtobool
from ..bg import remove from ..bg import remove
@ -90,7 +91,10 @@ def main():
full_paths += glob.glob(path + "/*") full_paths += glob.glob(path + "/*")
for fi in files: 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 continue
with open(fi, "rb") as input: with open(fi, "rb") as input: