添加image similarity

1 files changed, 71 insertions, 0 deletions

diff --git a/cv/image_similarity/cnn/feature_similarity.py b/cv/image_similarity/cnn/feature_similarity.py
new file mode 100644
index 0000000..ff10250
--- /dev/null
+++ b/cv/image_similarity/cnn/feature_similarity.py
@@ -0,0 +1,71 @@
+import random
+import torch
+from cv.image_similarity.cnn.build_models import IR_50, IR_101, IR_152
+import torchvision
+import os
+from PIL import Image
+
+
+to_torch_tensor = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),
+                                  torchvision.transforms.Normalize([0.5, 0.5, 0.5],[0.5, 0.5, 0.5])])
+
+
+
+def l2_norm(input, axis=1):
+    norm = torch.norm(input, 2, axis, True)
+    output = torch.div(input, norm)
+
+    return output
+
+
+def init_model(model, param, device):
+    m = model([112, 112])
+    m.eval()
+    m.to(device)
+    m.load_state_dict(torch.load(param, map_location=torch.device('cpu')))
+    return m
+
+
+def get_model_pool(device):
+    model_pool = []
+    # double
+    model_pool.append(init_model(IR_50, 'models/backbone_ir50_ms1m_epoch120.pth', device))
+    # model_pool.append(init_model(IR_50, 'models/backbone_ir50_ms1m_epoch120.pth', device))
+    #
+    # model_pool.append(init_model(IR_50, 'models/Backbone_IR_50_LFW.pth', device))
+    # model_pool.append(init_model(IR_101, 'models/Backbone_IR_101_Batch_108320.pth', device))
+    # model_pool.append(init_model(IR_152, 'models/Backbone_IR_152_MS1M_Epoch_112.pth', device))
+    return model_pool
+
+
+def get_model(device):
+    return init_model(IR_50, 'models/backbone_ir50_ms1m_epoch120.pth', device)
+
+
+device = torch.device('cpu')
+model = get_model(device)
+print('----models load over----')
+images = os.listdir('../imgs/lena')
+
+
+vectors = []
+
+for img in images:
+    print(img)
+    img = Image.open('../imgs/lena/' + img).convert('RGB')
+    # print(dir(img))
+    print(img.size, img.mode)
+    img = to_torch_tensor(img.resize((112, 112), Image.ANTIALIAS))
+    img = img.unsqueeze_(0).to(device)
+    feature = model(img)
+    vectors.append(l2_norm(feature).detach_())
+
+print('----vectors calculate over----')
+
+
+for i in range(len(vectors)):
+    for j in range(len(vectors)):
+        # consine distance
+        dist = (vectors[i]*vectors[j]).sum().item()
+        print(images[i], images[j], dist)
+    print('-------------')