사진첩에서 강아지를 찾아라!(AI torchvision활용, 파이썬)

안녕하세요.

 

최근에는 메타데이터부터 정형화된 데이터 분석위주로 했는데,

요즘 이미지 분석도 나중에 쓸 일이 있을거 같아 경험삼아

사진첩에서 강아지를 찾는 코드를 짜보았습니다!

(AI torchvision 활용)

 

1. 먼저 특정 폴더에 사진을 저장해둡니다.

 - 저는 pixabay의 무료 사진들을 가져왔습니다.(강아지, 고양이, 기린 사진 1개씩)

강아지 사진

 

고양이 사진

 

기린 사진

 

2. 파이썬으로 torchvision을 설치해봅니다.

 - pip install torch torchvision

 

3. torchvision을 import하여 ResNet-18 모델을 로드합니다.

 # 사전 훈련된 ResNet-18 모델 로드
model = models.resnet18(pretrained=True)
model.eval()  # 평가 모드로 설정

 

4. 강아지를 찾는 코드 완성

 - 마지막 코드를 참고하세요.

 - 강아지외의 객체를 찾는 방법은 ImageNet 클래스 인덱스를 변경하면서 찾아보세요~

5. 수행결과

 : 강아지를 잘 찾아냅니다.

강아지가 있는 이미지: ['shepherd-dog-4357790_1280.jpg']

강아지 사진

6. 전체 코드

import torch
import torchvision.transforms as transforms
from torchvision import models
from PIL import Image
import os

# 사전 훈련된 ResNet-18 모델 로드
model = models.resnet18(pretrained=True)
model.eval()  # 평가 모드로 설정

# 이미지를 모델에 맞는 형태로 변환하는 전처리 함수
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

def is_dog_in_image(image_path):
    """이미지 경로를 받아 해당 이미지에 강아지가 있는지 확인하는 함수"""
    img = Image.open(image_path).convert('RGB')
    img_t = preprocess(img)
    batch_t = torch.unsqueeze(img_t, 0)

    with torch.no_grad():
        output = model(batch_t)
    
    # ImageNet 클래스 인덱스 151 ~ 268은 강아지에 해당
    probabilities = torch.nn.functional.softmax(output, dim=1)[0] * 100
    _, indices = torch.sort(output, descending=True)
    return any(151 <= index and index <= 268 for index in indices[0][:5])

def find_dogs_in_folder(folder_path):
    """폴더 내의 모든 이미지를 검사하여 강아지가 있는 이미지를 찾는 함수"""
    dog_images = []
    # 이미지 파일 확장자 목록
    valid_extensions = ['.jpg', '.jpeg', '.png', '.gif']
    for image_name in os.listdir(folder_path):
        # 파일 확장자 확인
        if any(image_name.lower().endswith(ext) for ext in valid_extensions):
            image_path = os.path.join(folder_path, image_name)
            if is_dog_in_image(image_path):
                dog_images.append(image_name)
    return dog_images

# 특정 폴더 내에서 강아지 이미지 찾기
folder_path = r'C:\Users\funsc\Mywork\AI APP(Chat GPT)\torchvision'  # 이 경로는 실제 경로로 변경해야 합니다.
dog_images = find_dogs_in_folder(folder_path)
print("강아지가 있는 이미지:", dog_images)