[ChatGPT] 冷蔵庫内の写真から「おすすめレシピ」を受け取ってみました　〜食品は、Segment Anything と 転移学習した分類モデルで検出してます〜

import os
import math
import numpy as np
import cv2
import torch
import matplotlib.pyplot as plt
from torch.nn import functional as F
from torchvision import transforms
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
import japanize_matplotlib
import openai


# 分類モデル(resnet)
class ImageClassification:
    def __init__(self, device, model, label_file):
        print("init Image Classification")
        self.device = device

        self.model = torch.load(model)
        self.model.eval()
        self.model.to(self.device)

        with open(label_file, mode="r") as f:
            self.class_names = f.read().split(",")
        self.class_names.remove("")
        print(self.class_names)

    def inference(self, image):
        preprocess = transforms.Compose(
            [
                transforms.ToPILImage(),
                transforms.Resize(256),  # サイズ変更　(256, 256)
                transforms.CenterCrop(224),  # 中心で切り抜く　(224, 224)
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
                ),
            ]
        )
        input_tensor = preprocess(image)
        input_batch = input_tensor.unsqueeze(0)
        input_batch = input_batch.to(self.device)

        with torch.no_grad():
            outputs = self.model(input_batch)

        batch_probs = F.softmax(outputs, dim=1)
        batch_probs, batch_indices = batch_probs.sort(dim=1, descending=True)

        for probs, indices in zip(batch_probs, batch_indices):
            for k in range(1):
                return self.class_names[indices[k]], probs[k]


# SAM
class SegmentAnything:
    def __init__(self, device, model_type, sam_checkpoint):
        print("init Segment Anything")
        self.device = device
        self.sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
        self.sam.to(self.device)

    def generat_masks(self, image):
        print("generate masks by Segment Anything")
        # 検出するマスクの品質　（default: 0.88)
        # pred_iou_thresh = 0.96 640*480の時ちょうどよかった
        pred_iou_thresh = 0.98  #  960×720の時はこれぐらい
        mask_generator_2 = SamAutomaticMaskGenerator(
            model=self.sam, pred_iou_thresh=pred_iou_thresh
        )
        masks = mask_generator_2.generate(image)
        self.sorted_anns = sorted(masks, key=(lambda x: x["area"]), reverse=True)
        self._length = len(masks)
        print("{} masks generated.".format(self._length))

    @property
    def length(self):
        return self._length

    def get(self, index):
        pixels = self.sorted_anns[index]["area"]
        mask = np.expand_dims(self.sorted_anns[index]["segmentation"], axis=-1)
        return mask, pixels


def read_image(file_name):
    image = cv2.imread(file_name)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    height, width = image.shape[:2]
    return image, height, width


# GPT
class GPT3:
    def __init__(self):
        openai.api_key = os.environ["OPENAI_API_KEY"]
        print("init gpt-3.5-turbo")

    def recipe_recommend(self, food_list):
        prompt = self._create_prompt(food_list)
        try:
            response = openai.ChatCompletion.create(
                model="gpt-3.5-turbo", messages=prompt, temperature=0.7
            )
            print(response.choices[0]["message"]["content"].strip())
        except Exception as e:
            print("Exception:", e.args)

    def _create_prompt(self, food_list):
        list_str = ""
        for food in food_list:
            list_str += "{},".format(food)

        order = """
あなたは高級料理店の専門シェフです。 ユーザから与えられた食材のリスト使用して1種類のレシピを教えてくれます。
レシピには、キャッチーなタイトル、調理するのに必要な時間、また、レシピと同様に、リストに成分部分を含めます
"""
        role = "role"
        content = "content"
        system = "system"
        user = "user"
        assistant = "assistant"

        prompt = []
        prompt.append({role: system, content: order})
        prompt.append({role: user, content: "食材リストは、{} です".format(list_str)})
        prompt.append({role: assistant, content: ""})
        return prompt


# エリアを検出するクラス
class Area:
    def __init__(self, width, height):
        self.x1 = width
        self.x2 = 0
        self.y1 = height
        self.y2 = 0

    def set(self, x, y):
        if x < self.x1:
            self.x1 = x
        if self.x2 < x:
            self.x2 = x
        if y < self.y1:
            self.y1 = y
        if self.y2 < y:
            self.y2 = y

    def get(self):
        return self.x1, self.y1, self.x2, self.y2


# 検出したデータを表現するクラス
class Data:
    def __init__(self, image, contours, class_name, probs):
        self._image = image
        self._contours = contours
        self._class_name = class_name
        self._probs = probs

    @property
    def image(self):
        return self._image

    @property
    def contours(self):
        return self._contours

    @property
    def class_name(self):
        return self._class_name

    @property
    def probs(self):
        return self._probs


# 名前の索引
def get_name(class_name):
    if class_name == "sausage":
        return "ソーセージ"
    if class_name == "cabbage":
        return "キャベツ"
    if class_name == "potato":
        return "じゃがいも"
    if class_name == "milk":
        return "牛乳"
    if class_name == "pumpkin":
        return "かぼちゃ"
    if class_name == "green_pepper":
        return "ピーマン"
    if class_name == "nasubi":
        return "なすび"
    if class_name == "paprika":
        return "パプリカ"
    if class_name == "carrot":
        return "人参"
    if class_name == "apple":
        return "りんご"
    if class_name == "chikin":
        return "とり肉"
    if class_name == "pork":
        return "豚肉"
    if class_name == "tomato":
        return "トマト"
    if class_name == "banana":
        return "バナナ"
    if class_name == "onion":
        return "玉ねぎ"
    return "食品ではない".format(class_name)


def job(filename, sam, ic, gpt3):
    print("==========================================================")
    print("おすすめレシピ")
    print("Configured in gpt-3.5-turbo, Segment Anything and resnet. ")
    print("==========================================================")
    # 画像の読み込み
    image, height, width = read_image(filename)
    print("read_image {} width:{} height:{}".format(filename, width, height))
    print("この写真を元に、おすすめのレシピを提案します。")
    plt.imshow(image)
    plt.axis("off")
    plt.show()

    # マスク生成
    sam.generat_masks(image)

    data_list = []
    # 一定のサイズのものだけを抽出する
    max_pixels = 13000
    min_pixels = 2300

    for index in range(sam.length):
        mask, pixels = sam.get(index)

        # 一定範囲のサイズのみ対象にする
        if pixels < min_pixels or max_pixels < pixels:
            continue

        # 輪郭検出用の2値のテンポラリ画像
        mono_image = np.full(np.array([height, width, 1]), 255, dtype=np.uint8)
        # 個々の食品を切取るためのテンポラリ画像
        food_image = image.copy()

        area = Area(width, height)
        for y in range(height):
            for x in range(width):
                if mask[y][x]:
                    mono_image[y][x] = 0  # ２値画像は、マスク部分を黒にする
                    area.set(x, y)
                else:
                    food_image[y][x] = [255, 255, 255]  # 食品切取り画像は、マスク部分以外を白にする

        # 検出範囲
        x1, y1, x2, y2 = area.get()
        # 食品の輪郭を取得する
        contours, _ = cv2.findContours(
            mono_image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
        )
        # 食品画像を切取る
        food_image = food_image[y1:y2, x1:x2]
        # 分類モデルで食品名を推論する
        class_name, probs = ic.inference(food_image)

        # nodeで始まる名前は、食品では無い
        if class_name.startswith("none") == False:
            data_list.append(Data(food_image, contours, class_name, probs))
            print(
                "{} {:.2f}% [{},{},{},{}]".format(
                    get_name(class_name), probs, x1, y1, x2, y2
                )
            )

    # おすすめレシピ
    print("----------------------------------------------------------")
    food_list = []
    for data in data_list:
        food_list.append(get_name(data.class_name))

    gpt3.recipe_recommend(food_list)
    print("----------------------------------------------------------")

    # 元画像に輪郭を描画
    for data in data_list:
        cv2.drawContours(image, data.contours, -1, color=[0, 255, 255], thickness=3)
    plt.imshow(image)

    # 検出した食品画像を表示
    W = math.ceil(len(data_list) / 2)
    H = 2
    fig = plt.figure(figsize=(5, 5))

    for i, data in enumerate(data_list):
        ax1 = fig.add_subplot(H, W, i + 1)
        ax1.set_title(
            "{} {:.2f}".format(get_name(data.class_name), data.probs), fontsize=10
        )
        plt.imshow(data.image)

    plt.axis("off")
    plt.show()

    # 画面クリア
    os.system("clear")


def main():
    print("PyTorch version:", torch.__version__)
    print("CUDA is available:", torch.cuda.is_available())
    device = "cuda"

    # Resnet
    ic = ImageClassification(device, "model.pth", "label.txt")
    # Segment Anything
    sam = SegmentAnything(device, "vit_h", "sam_vit_h_4b8939.pth")
    # GPT3
    gpt3 = GPT3()

    file_list = [
        "image001.jpg",
        "image002.jpg",
        "image003.jpg",
        "image004.jpg",
        "image005.jpg",
        "image006.jpg",
    ]
    for file in file_list:
        job(file, sam, ic, gpt3)


if __name__ == "__main__":
    main()