update

ieemoo-ai-filtervideo.py

@@ -1,23 +1,37 @@
 from flask import request,Flask, jsonify
 from checkobject import get_object_location
 import numpy as np
-import cv2
+import cv2, os, requests, urllib, sys,json, time
+from datetime import datetime
 
 app = Flask(__name__)
+def _progress(block_num, block_size, total_size):
+    '''回调函数
+       @block_num: 已经下载的数据块
+       @block_size: 数据块的大小
+       @total_size: 远程文件的大小
+    '''
+    sys.stdout.write('\r>> Downloading %.1f%%' % (float(block_num * block_size) / float(total_size) * 100.0))
+    sys.stdout.flush()
+
 @app.route('/filtervideo', methods=['POST'])
 def filtervideo():
-    videourls = request.form.get('videoUrls')
-    videoid = request.form.get('id')
-    videopath = os.sep.join(['data', videoid+'.mp4'])
-    barcode = request.form.get('barcode')
-    videourls = videourls.split(',')
+    #videourls = request.form.get('videoUrls')
+    dicdatas = request.get_json()
     results = []
-    for name in videourls:
-        videos = requests.get(url)
-        videos.save(videopath)
+    for dicdata in dicdatas:
+        date = datetime.today()
+        videourl = dicdata["videoUrls"]  
+        videoid = dicdata['id']
+        videopath = os.sep.join(['data', videoid+'.mp4'])
+        barcode = dicdata['barcode']
+        
+        urllib.request.urlretrieve(videourl, videopath, _progress)
+
         flag = get_object_location(videopath)
-        results.append({'id':videoids, 'result':flags})
-    return result
+        results.append({'id':videoid, 'result':flag})
+    print('{} >>>>{}'.format((date),(results)))
+    return json.dumps(results)
             
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=8085)

tools/segpredict.py

@@ -0,0 +1,43 @@
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+import time
+import numpy as np
+import cv2
+import torch
+import segmentation_models_pytorch as smp
+import imageio
+
+def predict(best_model=None):
+    DATA_DIR = './data/CamVid/'
+    x_test_dir = os.path.join(DATA_DIR, 'abc')
+
+    img_test = cv2.imread('data/CamVid/abc/pic_unscan_front.jpg')
+    height = img_test.shape[0]
+    weight = img_test.shape[1]
+
+    ENCODER = 'resnet18'
+    ENCODER_WEIGHTS = 'imagenet'
+    CLASSES = ['front']
+    ACTIVATION = 'sigmoid' # could be None for logits or 'softmax2d' for multiclass segmentation
+    DEVICE = 'cuda'
+
+    # 加载最佳模型
+    best_model = torch.load('./module/best_model.pth')
+
+    for imgpath in os.listdir(x_test_dir):
+        img = cv2.imread(os.sep.join([x_test_dir, imgpath]))
+        img = cv2.resize(img, (512,512))
+        img = img.transpose((2,0,1))
+        T1 = time.time()
+        x_tensor = torch.from_numpy(img).float().to(DEVICE).unsqueeze(0)
+        pr_mask = best_model.predict(x_tensor)
+        T2 = time.time()
+        print('>>>> {}'.format(T2-T1))
+        pr_mask = (pr_mask.squeeze().cpu().numpy().round())
+        print('>>>>>>>>> pr_mask{}'.format(pr_mask.shape))
+        pr_mask = cv2.resize(pr_mask[0,:,:], (weight, height))
+        imageio.imwrite('f_test_out.png', pr_mask)
+
+if __name__ == '__main__':
+    predict()

tools/segpredict.py.bak

@@ -0,0 +1,180 @@
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+import time
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+import albumentations as albu
+import torch
+import segmentation_models_pytorch as smp
+from torch.utils.data import Dataset as BaseDataset
+import imageio
+
+
+# ---------------------------------------------------------------
+### Dataloader
+
+class Dataset(BaseDataset):
+    """CamVid数据集。进行图像读取，图像增强增强和图像预处理.
+
+    Args:
+        images_dir (str): 图像文件夹所在路径
+        masks_dir (str): 图像分割的标签图像所在路径
+        class_values (list): 用于图像分割的所有类别数
+        augmentation (albumentations.Compose): 数据传输管道
+        preprocessing (albumentations.Compose): 数据预处理
+    """
+    # CamVid数据集中用于图像分割的所有标签类别
+    #CLASSES = ['sky', 'building', 'pole', 'road', 'pavement',
+    #           'tree', 'signsymbol', 'fence', 'car',
+    #           'pedestrian', 'bicyclist', 'unlabelled']
+    CLASSES = ['front']
+
+    def __init__(
+            self,
+            images_dir,
+            # masks_dir,
+            classes=None,
+            augmentation=None,
+            preprocessing=None,
+    ):
+        self.ids = os.listdir(images_dir)
+        self.images_fps = [os.path.join(images_dir, image_id) for image_id in self.ids]
+
+        # convert str names to class values on masks
+        self.class_values = [self.CLASSES.index(cls.lower()) for cls in classes]
+
+        self.augmentation = augmentation
+        self.preprocessing = preprocessing
+
+    def __getitem__(self, i):
+
+        # read data
+        image = cv2.imread(self.images_fps[i])
+        image = cv2.resize(image, (512, 512))   # 改变图片分辨率
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        # 图像增强应用
+        if self.augmentation:
+            sample = self.augmentation(image=image)
+            image = sample['image']
+
+        # 图像预处理应用
+        if self.preprocessing:
+            sample = self.preprocessing(image=image)
+            image = sample['image']
+
+        return image
+
+    def __len__(self):
+        return len(self.ids)
+
+# ---------------------------------------------------------------
+
+def get_validation_augmentation():
+    """调整图像使得图片的分辨率长宽能被32整除"""
+    test_transform = [
+        albu.PadIfNeeded(384, 480)
+    ]
+    return albu.Compose(test_transform)
+
+
+def to_tensor(x, **kwargs):
+    return x.transpose(2, 0, 1).astype('float32')
+
+
+def get_preprocessing(preprocessing_fn):
+    """进行图像预处理操作
+
+    Args:
+        preprocessing_fn (callbale): 数据规范化的函数
+            (针对每种预训练的神经网络)
+    Return:
+        transform: albumentations.Compose
+    """
+
+    _transform = [
+        albu.Lambda(image=preprocessing_fn),
+        albu.Lambda(image=to_tensor),
+    ]
+    return albu.Compose(_transform)
+
+
+# 图像分割结果的可视化展示
+def visualize(**images):
+    """PLot images in one row."""
+    n = len(images)
+    plt.figure(figsize=(16, 5))
+    for i, (name, image) in enumerate(images.items()):
+        plt.subplot(1, n, i + 1)
+        plt.xticks([])
+        plt.yticks([])
+        plt.title(' '.join(name.split('_')).title())
+        plt.imshow(image)
+    plt.show()
+
+
+# ---------------------------------------------------------------
+if __name__ == '__main__':
+
+    DATA_DIR = './data/CamVid/'
+
+    x_test_dir = os.path.join(DATA_DIR, 'abc')
+
+    img_test = cv2.imread('data/CamVid/abc/pic_unscan_front.jpg')
+    height = img_test.shape[0]
+    weight = img_test.shape[1]
+
+    ENCODER = 'resnet18'
+    ENCODER_WEIGHTS = 'imagenet'
+    CLASSES = ['front']
+    ACTIVATION = 'sigmoid' # could be None for logits or 'softmax2d' for multiclass segmentation
+    DEVICE = 'cuda'
+
+    # 按照权重预训练的相同方法准备数据
+    preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)
+
+    # 加载最佳模型
+    best_model = torch.load('./module/best_model.pth')
+
+    # 创建检测数据集
+    predict_dataset = Dataset(
+        x_test_dir,
+        augmentation=get_validation_augmentation(),
+        preprocessing=get_preprocessing(preprocessing_fn),
+        classes=CLASSES,
+    )
+
+#    # 对检测图像进行图像分割并进行图像可视化展示
+#    predict_dataset_vis = Dataset(
+#        x_test_dir,
+#        classes=CLASSES,
+#    )
+
+    for i in range(len(predict_dataset)):
+        # 原始图像image_vis
+        #image_vis = predict_dataset_vis[i].astype('uint8')
+        image = predict_dataset[i]
+        print('>>>>> {}>>>size{}'.format(type(image), image.shape))
+
+        # 通过图像分割得到的0-1图像pr_mask
+        T1 = time.time()
+        x_tensor = torch.from_numpy(image).to(DEVICE).unsqueeze(0)
+        pr_mask = best_model.predict(x_tensor)
+        T2 = time.time()
+        print('>>>>>> {}'.format(T2-T1))
+        pr_mask = (pr_mask.squeeze().cpu().numpy().round())
+
+        # 恢复图片原来的分辨率
+        #image_vis = cv2.resize(image_vis, (weight, height))
+        #pr_mask = cv2.resize(pr_mask, (weight, height))
+        pr_mask = cv2.resize(pr_mask[0,:,:], (weight, height))
+        # 保存图像分割后的黑白结果图像
+        imageio.imwrite('f_test_out.png', pr_mask)
+        # 原始图像和图像分割结果的可视化展示
+#        visualize(
+#            image=image_vis,
+#            predicted_mask=pr_mask
+#        )
+