dataset_api

import json
import os
import requests
import base64
import cv2 as cv
import numpy as np
import base64
from tqdm import tqdm
class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(NpEncoder, self).default(obj)

def prepare_pic(pic_path):
    img_mat = cv.imdecode(np.fromfile(pic_path, dtype=np.uint8),  -1)
    img_data = cv.imencode(".jpg", img_mat)[1].tobytes()
    img_data = base64.b64encode(img_data)
    img_data = img_data.decode('utf-8')
    return img_data

def save_pic(img_return,save_pic_path):
    img_data = base64.b64decode(img_return)
    img_data = np.fromstring(img_data, np.int8)
    img = cv.imdecode(img_data, cv.IMREAD_COLOR)
    cv.imencode('.jpg', img)[1].tofile(save_pic_path)

def detect_images(img_dir, save_img_dir, argument='img', url=None):
    img_list = [x for x in os.listdir(img_dir) if x.endswith('.jpg')]
    for img in tqdm(img_list):
        img_path = os.path.join(img_dir, img)
        img_data = prepare_pic(img_path)
        if img_data is None:
            print('not find img!', img_data)
            continue
        data_obj = {'img': img_data, 'argument': argument}
        r = requests.post(url, json.dumps(data_obj))
        if r.status_code != 200:
            print(r.status_code,'error!')
            continue
        b = r.content.decode('utf-8')
        data_return = json.loads(b)
        img_name = os.path.split(img_path)[-1]
        save_img_path = os.path.join(save_img_dir, img_name)
        if data_obj['argument'] == 'img':
            save_pic(data_return['img'], save_img_path)
            print()
        else:
            text_path = save_img_path.replace('.jpg', '.json')
            with open(text_path, 'w') as f:
                json.dump(data_return, f)

def detect_image(img_path, save_img_dir, argument='img', url=None):
    img_data = prepare_pic(img_path)
    if img_data is None:
        print('not find img!', img_data)
        exit()
    data_obj = {'img': img_data, 'argument': argument}
    r = requests.post(url, json.dumps(data_obj))
    print('status_code:', r.status_code)
    b = r.content.decode('utf-8')
    data_return = json.loads(b)
    img_name = os.path.split(img_path)[-1]
    save_img_path = os.path.join(save_img_dir, img_name)
    if data_obj['argument'] == 'img':
        save_pic(data_return['img'], save_img_path)
    else:
        text_path = save_img_path.replace('.jpg', '.json')
        with open(text_path, 'w') as f:
            json.dump(data_return, f)


if __name__ == '__main__':
    url = 'http://0.0.0.0:2750/yolact/detect'  # GPUA
    img_dir = r'F:\data_sets\street_dataset2\test_class'
    save_img_dir = r'C:\Users\K0802363\Desktop\BIAOZHUHOU'
    detect_images(img_dir, save_img_dir,url)