diff --git a/.checkobject.py.swp b/.checkobject.py.swp new file mode 100644 index 0000000..ec3809f Binary files /dev/null and b/.checkobject.py.swp differ diff --git a/.ieemoo-ai-filtervideo.py.swp b/.ieemoo-ai-filtervideo.py.swp new file mode 100644 index 0000000..2fd4904 Binary files /dev/null and b/.ieemoo-ai-filtervideo.py.swp differ diff --git a/data/1.mp4 b/data/1.mp4 new file mode 100644 index 0000000..73d5240 Binary files /dev/null and b/data/1.mp4 differ diff --git a/ieemoo-ai-filtervideo.py b/ieemoo-ai-filtervideo.py index 29acc38..50f9986 100644 --- a/ieemoo-ai-filtervideo.py +++ b/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) diff --git a/tools/segpredict.py b/tools/segpredict.py new file mode 100644 index 0000000..85e409f --- /dev/null +++ b/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() diff --git a/tools/segpredict.py.bak b/tools/segpredict.py.bak new file mode 100644 index 0000000..403370c --- /dev/null +++ b/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 +# ) +