|
- from __future__ import absolute_import
- from __future__ import division
- from __future__ import print_function
- from __future__ import unicode_literals
- import pickle
- import os
- import sys
- #sys.path.append(r"/home/xieliang/example/django_alg_service/alg_manager/CBN-archived-002/CBN-archived-codes")
- import random
- import tqdm
- import time
- import numpy as np
- import torch
- from torch.utils.data import DataLoader
- from PIL import Image
- from config import opt
- from io_stream import data_manager
-
- from frameworks.models import ResNetBuilder
- from frameworks.evaluating import evaluator_manager
-
- sys.path.append(r"/home/xieliang/example/django_alg_service/alg_manager/CBN-archived-002/CBN-archived-codes")
-
- from utils.serialization import Logger, load_previous_model, load_moco_model
- from utils.transforms import TestTransform
- import os
- from urllib import request
- import base64
-
- torch.set_printoptions(profile="full")
- import cv2
- from sklearn.metrics.pairwise import cosine_similarity
-
- class CBN_FeatureExtraction(object):
-
- def __init__(self):
-
- model = ResNetBuilder()
- model_path = os.path.join("./pytorch-ckpt/current","market_demo",
- 'model_best.pth.tar')
- model = load_previous_model(model, model_path, load_fc_layers=False)
- model.eval()
- #self.reid_evaluator = evaluator_manager.init_evaluator(testset_name, model, flip=True)
-
- #self.dataset = data_manager.init_dataset(name=testset_name,
- # num_bn_sample=320)\
- self.model = model
- self.trans = TestTransform(opt.height,opt.width)
-
- def do_alg(self, input):
- '''
- Arg
- ---
- input: A dict, key "image" : value: numpy.array.
-
-
- Return
- ---
- person feature.
- size: 2048 dim
- type: numpy.array
- '''
- img = Image.fromarray(input["image"][0])
- img = self.trans(img).view(1,3,256,128)
-
- with torch.no_grad():
- f,_ = self.model(img)
- f = f.cpu().numpy().squeeze()
- return {"result":f.reshape(1,-1)}
-
-
-
-
- if __name__ == '__main__':
- featExtractor = CBN_FeatureExtraction()
- imgs = [
- "http://0.0.0.0:801/0089_c6s1_015526_03.jpg",
- "http://0.0.0.0:801/0091_c1s1_021826_02.jpg",
- "http://0.0.0.0:801/0089_c6s1_013551_02.jpg"
- ]
- input_q_url = imgs[0]
- input_g_url = imgs[1]
- resp = request.urlopen(input_q_url)
- image = np.asarray(bytearray(resp.read()), dtype="uint8")
- image = cv2.imdecode(image, cv2.IMREAD_COLOR)
- img = image.astype(np.uint8)
- img = img[np.newaxis, ...]
-
- feat_q = featExtractor.do_alg({"image": img})
- feat_q = feat_q["result"]
- # gallery
- resp1 = request.urlopen(input_g_url)
- image1 = np.asarray(bytearray(resp1.read()), dtype="uint8")
- image1 = cv2.imdecode(image1, cv2.IMREAD_COLOR)
- img1 = image1.astype(np.uint8)
- img1 = img1[np.newaxis, ...]
- feat_g = featExtractor.do_alg({"image": img1})
- feat_g = feat_g["result"]
- cos_similary = cosine_similarity(feat_q, feat_g)
- print("Different person:",cos_similary)
- input_g_url = imgs[2]
- resp1 = request.urlopen(input_g_url)
- image1 = np.asarray(bytearray(resp1.read()), dtype="uint8")
- image1 = cv2.imdecode(image1, cv2.IMREAD_COLOR)
- img1 = image1.astype(np.uint8)
- img1 = img1[np.newaxis, ...]
- feat_g = featExtractor.do_alg({"image": img1})
- feat_g = feat_g["result"]
- cos_similary = cosine_similarity(feat_q, feat_g)
- print("Same person:",cos_similary)
|