Add files via upload

Author: skx6

image_retrieval.py

@@ -0,0 +1,89 @@
+import os
+import cv2
+import time
+from datetime import timedelta
+from retrieval.create_thumb_images import create_thumb_images
+from flask import Flask, render_template, request, redirect, url_for, make_response,jsonify, flash
+from retrieval.retrieval import load_model, load_data, extract_feature, load_query_image, sort_img, extract_feature_query
+
+# Create thumb images.
+create_thumb_images(full_folder='./static/image_database/',
+                    thumb_folder='./static/thumb_images/',
+                    suffix='',
+                    height=200,
+                    del_former_thumb=True,
+                    )
+
+# Prepare data set.
+data_loader = load_data(data_path='./static/image_database/',
+                        batch_size=2,
+                        shuffle=False,
+                        transform='default',
+                        )
+
+# Prepare model.
+model = load_model(pretrained_model='./retrieval/models/net_best.pth', use_gpu=True)
+
+# Extract database features.
+gallery_feature, image_paths = extract_feature(model=model, dataloaders=data_loader)
+
+# Picture extension supported.
+ALLOWED_EXTENSIONS = set(['png', 'jpg', 'JPG', 'PNG', 'bmp', 'jpeg', 'JPEG'])
+def allowed_file(filename):
+    return '.' in filename and filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS
+
+app = Flask(__name__)
+# Set static file cache expiration time
+# app.send_file_max_age_default = timedelta(seconds=1)
+
+
+@app.route('/', methods=['POST', 'GET'])  # add route
+def image_retrieval():
+
+    basepath = os.path.dirname(__file__)    # current path
+    upload_path = os.path.join(basepath, 'static/upload_image','query.jpg')
+
+    if request.method == 'POST':
+        if request.form['submit'] == 'upload':
+            if len(request.files) == 0:
+                return render_template('upload_finish.html', message='Please select a picture file!')
+            else:
+                f = request.files['picture']
+         
+                if not (f and allowed_file(f.filename)):
+                    # return jsonify({"error": 1001, "msg": "Examine picture extension, only png, PNG, jpg, JPG, or bmp supported."})
+                    return render_template('upload_finish.html', message='Examine picture extension, png、PNG、jpg、JPG、bmp support.')
+                else:
+
+                    f.save(upload_path)
+             
+                    # transform image format and name with opencv.
+                    img = cv2.imread(upload_path)
+                    cv2.imwrite(os.path.join(basepath, 'static/upload_image', 'query.jpg'), img)
+             
+                    return render_template('upload_finish.html', message='Upload successfully!')
+
+        elif request.form['submit'] == 'retrieval':
+            start_time = time.time()
+            # Query.
+            query_image = load_query_image('./static/upload_image/query.jpg')
+            # Extract query features.
+            query_feature = extract_feature_query(model=model, img=query_image)
+            # Sort.
+            similarity, index = sort_img(query_feature, gallery_feature)
+            sorted_paths = [image_paths[i] for i in index]
+
+            print(sorted_paths)
+            tmb_images = ['./static/thumb_images/' + os.path.split(sorted_path)[1] for sorted_path in sorted_paths]
+            # sorted_files = [os.path.split(sorted_path)[1] for sorted_path in sorted_paths]
+
+            return render_template('retrieval.html', message="Retrieval finished, cost {:3f} seconds.".format(time.time() - start_time),
+            	sml1=similarity[0], sml2=similarity[1], sml3=similarity[2], sml4=similarity[3], sml5=similarity[4], sml6=similarity[5],
+            	img1_tmb=tmb_images[0], img2_tmb=tmb_images[1],img3_tmb=tmb_images[2],img4_tmb=tmb_images[3],img5_tmb=tmb_images[4],img6_tmb=tmb_images[5])
+
+    return render_template('upload.html')
+
+
+if __name__ == '__main__':
+    # app.debug = True
+    app.run(host='0.0.0.0', port=8080, debug=True)

retrieval/create_thumb_images.py

@@ -0,0 +1,38 @@
+import os
+import cv2
+
+
+def del_file(path):
+    for i in os.listdir(path):
+        path_file = os.path.join(path, i)
+        if os.path.isfile(path_file):
+            os.remove(path_file)
+        else:
+            del_file(path_file)
+
+
+def create_thumb_images(full_folder, thumb_folder, suffix='thumb', height=100, del_former_thumb=False):
+    if del_former_thumb:
+        del_file(thumb_folder)
+    for image_file in os.listdir(full_folder):
+        image = cv2.imread(full_folder + image_file)
+        height_src, width_src, _ = image.shape
+        #print('width: {}, height: {}'.format(width_src, height_src))
+
+        width = (height / height_src) * width_src
+        # print(' Thumb width: {}, height: {}'.format(width, height))
+
+        resized_image = cv2.resize(image, (int(width), int(height)))
+
+        image_name, image_extension = os.path.splitext(image_file)
+        cv2.imwrite(thumb_folder + image_name + suffix + image_extension, resized_image)
+    print('Creating thumb images finished.')
+
+
+if __name__ == '__main__':
+    create_thumb_images(full_folder='./save_pic/',
+                        thumb_folder='./thumb_images/',
+                        suffix='',
+                        height=200,
+                        del_former_thumb=True,
+                        )

retrieval/retrieval.py

@@ -0,0 +1,191 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import print_function, division
+
+import argparse
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+from torchvision import datasets, models, transforms
+import os
+
+
+from torch.utils.data import dataloader, Dataset
+from PIL import Image
+
+
+def get_file_list(file_path_list, sort=True):
+    """
+    Get list of file paths in one folder.
+    :param file_path: A folder path or path list.
+    :return: file list: File path list of
+    """
+    import random
+    if isinstance(file_path_list, str):
+        file_path_list = [file_path_list]
+    file_lists = []
+    for file_path in file_path_list:
+        assert os.path.isdir(file_path)
+        file_list = os.listdir(file_path)
+        if sort:
+            file_list.sort()
+        else:
+            random.shuffle(file_list)
+        file_list = [file_path + file for file in file_list]
+        file_lists.append(file_list)
+    if len(file_lists) == 1:
+        file_lists = file_lists[0]
+    return file_lists
+
+
+class Gallery(Dataset):
+    """
+    Images in database.
+    """
+
+    def __init__(self, image_paths, transform=None):
+        super().__init__()
+
+        self.image_paths = image_paths
+        self.transform = transform
+
+    def __getitem__(self, index):
+        image_path = self.image_paths[index]
+        image = Image.open(image_path).convert('RGB')
+
+        if self.transform is not None:
+            image = self.transform(image)
+
+        return image, image_path
+
+    def __len__(self):
+        return len(self.image_paths)
+
+
+def load_data(data_path, batch_size=1, shuffle=False, transform='default'):
+    data_transform = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ]) if transform == 'default' else transform
+
+    image_path_list = get_file_list(data_path)
+
+    gallery_data = Gallery(image_paths=image_path_list,
+                           transform=data_transform,
+                           )
+
+    data_loader = dataloader.DataLoader(dataset=gallery_data,
+                                        batch_size=batch_size,
+                                        shuffle=shuffle,
+                                        num_workers=0,
+                                        )
+    return data_loader
+
+
+def extract_feature(model, dataloaders, use_gpu=True):
+    features = torch.FloatTensor()
+    path_list = []
+
+    use_gpu = use_gpu and torch.cuda.is_available()
+    for img, path in dataloaders:
+        img = img.cuda() if use_gpu else img
+        input_img = Variable(img.cuda())
+        outputs = model(input_img)
+        ff = outputs.data.cpu()
+        # norm feature
+        fnorm = torch.norm(ff, p=2, dim=1, keepdim=True)
+        ff = ff.div(fnorm.expand_as(ff))
+        features = torch.cat((features, ff), 0)
+        path_list += list(path)
+    return features, path_list
+
+
+def extract_feature_query(model, img, use_gpu=True):
+    c, h, w = img.size()
+    img = img.view(-1,c,h,w)
+    use_gpu = use_gpu and torch.cuda.is_available()
+    img = img.cuda() if use_gpu else img
+    input_img = Variable(img)
+    outputs = model(input_img)
+    ff = outputs.data.cpu()
+    fnorm = torch.norm(ff,p=2,dim=1, keepdim=True)
+    ff = ff.div(fnorm.expand_as(ff))
+    return ff
+
+
+def load_query_image(query_path):
+    data_transforms = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+    query_image = datasets.folder.default_loader(query_path)
+    query_image = data_transforms(query_image)
+    return query_image
+
+
+def load_model(pretrained_model=None, use_gpu=True):
+    """
+
+    :param check_point: Pretrained model path.
+    :return:
+    """
+    model = models.resnet50(pretrained=False)
+    num_ftrs = model.fc.in_features
+    add_block = []
+    add_block += [nn.Linear(num_ftrs, 30)]  #number of training classes
+    model.fc = nn.Sequential(*add_block)
+    model.load_state_dict(torch.load(pretrained_model))
+
+    # remove the final fc layer
+    model.fc = nn.Sequential()
+    # change to test modal
+    model = model.eval()
+    use_gpu = use_gpu and torch.cuda.is_available()
+    if use_gpu:
+        model = model.cuda()
+    return model
+
+
+# sort the images
+def sort_img(qf, gf):
+    score = gf*qf
+    score = score.sum(1)
+    # predict index
+    s, index = score.sort(dim=0, descending=True)
+    s = s.cpu().data.numpy()
+    import numpy as np
+    s = np.around(s, 3)
+    return s, index
+
+
+if __name__ == '__main__':
+
+    # Prepare data.
+    data_loader = load_data(data_path='./test_pytorch/gallery/images/',
+                            batch_size=2,
+                            shuffle=False,
+                            transform='default',
+                            )
+
+    # Prepare model.
+    model = load_model(pretrained_model='./model/ft_ResNet50/net_best.pth', use_gpu=True)
+
+    # Extract database features.
+    gallery_feature, image_paths = extract_feature(model=model, dataloaders=data_loader)
+
+    # Query.
+    query_image = load_query_image('./test_pytorch/query/query.jpg')
+
+    # Extract query features.
+    query_feature = extract_feature_query(model=model, img=query_image)
+
+    # Sort.
+    similarity, index = sort_img(query_feature, gallery_feature)
+
+    sorted_paths = [image_paths[i] for i in index]
+    print(sorted_paths)
+

static/image_database/10.jpg

@@ -0,0 +1,42 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <title>Image Retrieval</title>
+</head>
+<body>
+    <div style="padding-left: 60px;">
+        <h1>Image Retrieval</h1>
+
+            <div style="width:450px; float:left; background:#FFFFFF">
+                <form action="" enctype='multipart/form-data' method='POST'>
+                    <input type="file" name="picture" style="margin-top:20px;"/>
+                    <br>
+                    <i>Upload your picture, please click here: </i>
+                    <input type="submit" value="upload" name="submit" class="button-new" style="margin-top:15px;"/>
+                    <br>
+                    <i>Start retrieval, please click here: </i>
+                    <input type="submit" value="retrieval" name="submit" class="button-new" style="margin-top:15px;"/>
+                </form>
+                <h3>{{message}}</h3>
+                <img src="{{ url_for('static', filename= './upload_image/query.jpg') }}" width="400" />
+            </div> 
+
+            <div style=" margin-left:450px; background:#FFFFFF">
+                    <h3>Retrieval result:</h3>
+                    <h4>Similarity of top 3: {{sml1}}, {{sml2}}, {{sml3}}</h4>
+                    <div>
+                        <img src="{{img1_tmb}}" height="200px" overflow=hidden/>
+                        <img src="{{img2_tmb}}" height="200px" overflow=hidden/>
+                        <img src="{{img3_tmb}}" height="200px" overflow=hidden/>
+                    </div>
+                    <h4>Similarity of 4 ~ 6: {{sml4}}, {{sml5}}, {{sml6}}</h4>
+                    <div>
+                        <img src="{{img4_tmb}}" height="200px" overflow=hidden/>
+                        <img src="{{img5_tmb}}" height="200px" overflow=hidden/>
+                        <img src="{{img6_tmb}}" height="200px" overflow=hidden/>
+                    </div>
+            </div>  
+    </div>
+</body>
+</html>

static/image_database/11.jpg

@@ -0,0 +1,21 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
+    <title>Image Retrieval</title>
+</head>
+<body>
+    <div style="padding-left: 60px;">
+        <h1>Image Retrieval</h1>
+        <form action="" enctype='multipart/form-data' method='POST'>
+            <input type="file" name="picture" style="margin-top:20px;"/>
+            <br>
+            <i>Upload your picture, please click here: </i>
+            <input type="submit" value="upload" name="submit" class="button-new" style="margin-top:15px;"/>
+            <br>
+            <i>Start retrieval, please click here: </i>
+            <input type="submit" value="retrieval" name="submit" class="button-new" style="margin-top:15px;"/>
+        </form>
+    </div>
+</body>
+</html>