新人请教一下怎么添加判断

[toyst]

原代码：

def extract_features(img_dir, model):
    try:
        cache = Cache('./tmp')
        feats = []
        names = []
        img_list = get_imgs(img_dir)
        total = len(img_list)
        cache['total'] = total
        for i, img_path in enumerate(img_list):
            try:
                norm_feat = model.resnet50_extract_feat(img_path)
                feats.append(norm_feat)
                names.append(img_path.encode())
                cache['current'] = i + 1
                print(f"Extracting feature from image No. {i + 1} , {total} images in total")
            except Exception as e:
                LOGGER.error(f"Error with extracting feature from image {e}")
                continue
        return feats, names
    except Exception as e:
        LOGGER.error(f"Error with extracting feature from image {e}")
        sys.exit(1)

大概位置是下面这里：

def extract_features(img_dir, model):
    try:
        cache = Cache('./tmp')
        feats = []
        names = []
        img_list = get_imgs(img_dir)
        total = len(img_list)
        cache['total'] = total
        #check if img_dir is a image or image_path
        #if it is image, create embedding and return, else, traverse the directory
        for i, img_path in enumerate(img_list):
            #extract feature
        return feats, names
    except Exception as e:
        LOGGER.error(f"Error with extracting feature from image {e}")
        sys.exit(1)

请问一下应该如何写啊？ 判断传入的是图片文件名还是目录

[corningsun]

filetype ?

[Latin]

import os
flag = os.path.exists(dir_name)
from pathlib import Path
filePath = path("/data/mydir")
flag = True if filePath.is_dir() else False

[advancejar]

你是刚学吗？
str.endswith('.jpg')
还有就是
os.path.isdir
这个函数最准，但是有坑，千万多测试

[deplivesb]

看了下楼主的发帖记录。。。。

[ChrisFreeMan]

@deplivesb 面向 V2EX 编程

[efaun]

@deplivesb #4
@ChrisFreeMan #5
哈哈哈, 要不你们建个知识星球把楼主拉进去吧

[imn1]

文件系统的路径问题，建议 os.scandir 和 pathlib 比较省事，基本上就是特殊的字符串处理
如果要判断真实图片是需要导入内容的，例如文件头判断，就是 file signature