feat: initial commit

.gitattributes

@@ -29,3 +29,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.ann filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+.DS_Store
+__pycache__
+input_img.png
+src/data

api.py

@@ -0,0 +1,101 @@
+import annoy
+from typing import Tuple
+import json
+
+EMBEDDING_DIMENSION = 4096
+ANNOY_INDEX_FILE = "models/index.ann"
+ANNOY_MAPPING_FILE = "models/mappings.json"
+IMG_RESIZE_SIZE = 224
+
+####
+
+def load_annoy_index(
+    index_file=ANNOY_INDEX_FILE,
+    mapping_file=ANNOY_MAPPING_FILE,
+) -> Tuple[annoy.AnnoyIndex, dict]:
+    """Load annoy index and associated mapping file"""
+
+    annoy_index = annoy.AnnoyIndex(f=EMBEDDING_DIMENSION, metric='euclidean')
+    annoy_index.load(index_file)
+
+    with open(ANNOY_MAPPING_FILE) as f:
+        mappings = json.load(f)
+
+    with open(mapping_file) as f:
+        mapping = json.load(f)
+        mapping = {int(k): v for k, v in mapping.items()}
+    return annoy_index, mappings
+
+###
+
+import torch
+from torch import optim, nn
+from torchvision import models, transforms
+model = models.vgg16(pretrained=True)
+
+
+
+# Transform the image, so it becomes readable with the model
+transform = transforms.Compose([
+  transforms.ToPILImage(),
+  # transforms.CenterCrop(512),
+  # transforms.Resize(448),
+  transforms.Resize((IMG_RESIZE_SIZE, IMG_RESIZE_SIZE)),
+  transforms.ToTensor()                              
+])
+
+import cv2
+
+class FeatureExtractor(nn.Module):
+  def __init__(self, model):
+    super(FeatureExtractor, self).__init__()
+		# Extract VGG-16 Feature Layers
+    self.features = list(model.features)
+    self.features = nn.Sequential(*self.features)
+		# Extract VGG-16 Average Pooling Layer
+    self.pooling = model.avgpool
+		# Convert the image into one-dimensional vector
+    self.flatten = nn.Flatten()
+		# Extract the first part of fully-connected layer from VGG16
+    self.fc = model.classifier[0]
+  
+  def forward(self, x):
+		# It will take the input 'x' until it returns the feature vector called 'out'
+    out = self.features(x)
+    out = self.pooling(out)
+    out = self.flatten(out)
+    out = self.fc(out) 
+    return out 
+
+# Initialize the model
+model = models.vgg16(pretrained=True)
+new_model = FeatureExtractor(model)
+
+# Change the device to GPU
+device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
+new_model = new_model.to(device)
+
+import PIL
+
+def analyze_image(
+    image, annoy_index, n_matches: int = 1, num_jitters: int = 1, model: str = "large"
+):
+    PIL.Image.fromarray(image).save("input_img.png")
+    img = cv2.imread("input_img.png")
+    # Transform the image
+    img = transform(img)
+    # Reshape the image. PyTorch model reads 4-dimensional tensor
+    # [batch_size, channels, width, height]
+    # img = img.reshape(1, 3, 448, 448)
+    img = img.reshape(1, 3, IMG_RESIZE_SIZE, IMG_RESIZE_SIZE)
+    img = img.to(device)
+    # We only extract features, so we don't need gradient
+    with torch.no_grad():
+        # Extract the feature from the image
+        feature = new_model(img)
+    # Convert to NumPy Array, Reshape it, and save it to features variable
+    v = feature.cpu().detach().numpy().reshape(-1)
+
+    results = annoy_index.get_nns_by_vector(v, n_matches, include_distances=True)
+
+    return results

app.py

@@ -0,0 +1,56 @@
+import gradio as gr
+
+import numpy as np
+import re
+from pathlib import Path
+from api import load_annoy_index, analyze_image
+
+annoy_index, mappings = load_annoy_index()
+
+def get_article_text():
+    article = Path("README.md").read_text()
+    # Remove the HuggingFace Space app information from the README
+    article = re.sub(r"^---.+---\s+", "", article, flags=re.MULTILINE + re.DOTALL)
+    return article
+
+
+
+def find_matching_images(input_img, n_matches: int = 10):
+    results = analyze_image(input_img, annoy_index, n_matches=n_matches)
+    
+
+    indexes = results[0]
+    # scores = results[1]
+
+    images = []
+
+    for i in range(len(indexes)):
+        index = str(indexes[i])
+
+        mapping = mappings[index]
+
+        url = mapping["url"]
+        if url != "":
+          images.append(url)
+
+    return images
+
+iface = gr.Interface(
+    find_matching_images,
+    title="類似貼り込み資料検索",
+    description="""類似する貼り込み資料を検索します。 Upload a picture and find out!
+    Give it a shot or try one of the sample images below.
+    
+    Built with ❤️ using great open-source libraries such as PyTorch and Annoy.""",
+    article=get_article_text(),
+    inputs=[
+        gr.inputs.Image(shape=None, label="Your image"),
+    ],
+    outputs=gr.Gallery(label="類似する貼り込み資料"),
+    examples=[
+        ["images/025_0085.jpg"],
+        ["images/046_0051.jpg"],
+    ],
+)
+
+iface.launch()

models/index.ann

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:008f512c1d862a9162807170d3ee5445df1bff76e26cb48f01ffabe3eb07284c
+size 66534800

src/build.py

@@ -0,0 +1,109 @@
+import torch
+from torch import optim, nn
+from torchvision import models, transforms
+from torchvision.models.vgg import VGG16_Weights
+
+
+class FeatureExtractor(nn.Module):
+  def __init__(self, model):
+    super(FeatureExtractor, self).__init__()
+		# Extract VGG-16 Feature Layers
+    self.features = list(model.features)
+    self.features = nn.Sequential(*self.features)
+		# Extract VGG-16 Average Pooling Layer
+    self.pooling = model.avgpool
+		# Convert the image into one-dimensional vector
+    self.flatten = nn.Flatten()
+		# Extract the first part of fully-connected layer from VGG16
+    self.fc = model.classifier[0]
+  
+  def forward(self, x):
+		# It will take the input 'x' until it returns the feature vector called 'out'
+    out = self.features(x)
+    out = self.pooling(out)
+    out = self.flatten(out)
+    out = self.fc(out) 
+    return out 
+
+# Initialize the model
+model = models.vgg16(weights=VGG16_Weights.DEFAULT)
+new_model = FeatureExtractor(model)
+
+# Change the device to GPU
+device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
+new_model = new_model.to(device)
+
+IMG_RESIZE_SIZE = 224
+IMG_PATH = "data"
+
+import cv2
+from tqdm import tqdm
+import numpy as np
+
+# Transform the image, so it becomes readable with the model
+transform = transforms.Compose([
+  transforms.ToPILImage(),
+  transforms.Resize((IMG_RESIZE_SIZE, IMG_RESIZE_SIZE)),
+  transforms.ToTensor()                              
+])
+
+# Will contain the feature
+features = []
+
+mappings = {}
+
+import glob
+# files = glob.glob("/content/drive/Shareddrives/ndl/kao/dataset 3/*.jpg")
+files = glob.glob(f"{IMG_PATH}/*.jpg")
+
+files.sort()
+
+for index in tqdm(range(len(files))):
+
+  path = files[index]
+
+  img = cv2.imread(path)
+  # Transform the image
+  img = transform(img)
+  # Reshape the image. PyTorch model reads 4-dimensional tensor
+  # [batch_size, channels, width, height]
+  # img = img.reshape(1, 3, 448, 448)
+  img = img.reshape(1, 3, IMG_RESIZE_SIZE, IMG_RESIZE_SIZE)
+  img = img.to(device)
+  # We only extract features, so we don't need gradient
+  with torch.no_grad():
+    # Extract the feature from the image
+    feature = new_model(img)
+  # Convert to NumPy Array, Reshape it, and save it to features variable
+  features.append(feature.cpu().detach().numpy().reshape(-1))
+
+  mappings[index] = {
+      "nconst": path.split("/")[-1].split(".")[0],
+      "name": "",
+      "url": ""
+  }
+
+# Convert to NumPy Array
+features = np.array(features)
+
+import json
+
+with open('mappings.json', mode='wt', encoding='utf-8') as file:
+  json.dump(mappings, file, ensure_ascii=False, indent=2)
+
+N_TREES = 1000
+
+from annoy import AnnoyIndex
+
+annoy_index = AnnoyIndex(features.shape[1], metric='euclidean')
+
+for i in range(len(features)):
+
+    feature = features[i]
+
+    annoy_index.add_item(i, feature)
+
+# k-d tree をビルドする
+annoy_index.build(n_trees=N_TREES)
+
+annoy_index.save("../models/index.ann")

src/dwn.py

@@ -0,0 +1,21 @@
+import json
+from urllib import request
+from tqdm import tqdm
+import os
+
+with open("/Users/nakamura/git/kunshujo/kunshujo/static/data/index.json") as f:
+    items = json.load(f)
+    for item in tqdm(items):
+        id = item["objectID"]
+        thumbnail = item["thumbnail"]
+        url = thumbnail.replace("/,300/", "/!512,512/")
+
+        opath = f"data/{id}.jpg"
+
+        if os.path.exists(opath):
+            continue
+
+        try:
+            request.urlretrieve(url, opath)
+        except:
+            pass

src/mapping.py

@@ -0,0 +1,27 @@
+import json
+
+with open('mappings.json') as f:
+    df = json.load(f)
+
+conf = {}
+
+with open("/Users/nakamura/git/kunshujo/kunshujo/static/data/index.json") as f:
+    items = json.load(f)
+    for item in items:
+        conf[item["objectID"]] = item        
+
+
+for index in df:
+    item = df[index]
+    id = item["nconst"]
+
+    if id not in conf:
+        continue
+
+    c = conf[id]
+
+    item["url"] = c["thumbnail"]
+    item["name"] = c["label"]
+
+with open('../models/mappings.json', mode='wt', encoding='utf-8') as file:
+  json.dump(df, file, ensure_ascii=False, indent=2)