Save code

# encoding: utf-8

import requests
import urllib.request
import numpy as np
import os ,stat
import pickle as pkl
import cv2 as cv
import xlrd

def getImg(url, basedir, filename):

    try:
        urllib.request.urlretrieve(url,filename=basedir+"{name}.jpg".format(name = filename))
    except IOError as e:
        print("IOE ERROR")
    except Exception as e:
        print("Exception")


def loadPkl(name):
    if not os.path.exists(name):
        color_feature = {}
        data = open(name, 'wb')
        pkl.dump(color_feature, data)
        data.close()
    data = open(name, 'rb')
    return pkl.load(data)


def savePkl(color_feature, name):
    data = open(name, 'wb')
    pkl.dump(color_feature, data)
    data.close()


def getImgInfoFromCsv(filename):
    data = xlrd.open_workbook(filename)
    try:
        sh = data.sheet_by_name("Sheet1")
    except:
        print(" There is no such document ！！！")
    nrows = sh.nrows  #  Get the number of lines 
    ncols = sh.ncols  #  Get the number of columns 
    res = {}  #  Get the data from the table and store it in the rowlist in 
    for i in range(1, nrows):
        img_id = sh.cell_value(i, 0)
        img_url = sh.cell_value(i, 1)
        res[img_id] = img_url
    return res


def getImgInfoFromTxt(filename):
    res = {}
    with open(filename, 'r', encoding='utf-8') as f:
        data = f.readlines()
    for i in data:
        if i.strip().split(" ")[0] not in res:
            res[i.strip().split(" ")[0]] = i.strip().split(" ")[1]
    return res



def getColorFeature(img):

    return 1, 2, 3, 4


# img_info example: {img_id, : "http://ww.png"}
def processColorFeature(img_info, color_feature):

    for i in img_info.keys():
        url = img_info[i]
        file = requests.get(url)
        img = cv.imdecode(np.fromstring(file.content, np.uint8), 1)

        a, b, c, d = getColorFeature(img)
        color_feature[i] = [a, b, c, d]

    return color_feature

def pklToTxt(source, dim):
    data = open(source, 'w')
    feature_info = pkl.load(data)
    with open(dim, 'r', encoding='utf-8') as f:
        for i in feature_info.keys():
            temp = i + " " + " ".join(feature_info[i]) + "\n"
            f.write(temp)


if __name__ == "__main__":
    savename = "banner_color.pickle"
    info_name = ""
    #get history color feature
    color_feature = loadPkl(savename)
    #get origin img information
    img_info = getImgInfoFromCsv(info_name)
    #process
    feature = processColorFeature(img_info, color_feature)
    #save result
    savePkl(feature, savename)