GitHub - video2pdfslidesvideo2pdfslides 库主要功能：从视频中提取图片，并去重后，转为PDF.1. 依赖与安装依赖项：imutils==0.5.4 opencv_python==4.5.2.52 img2pdf==0.4.1安装：git clone https://github.com/kaushikj/video2pdfslides cd video2pdfslides pip install -r requirements.txt运行：python video2pdfslides.py <video_path>2. video2pdfslideshttps://github.com/kaushikj/video2pdfslides/blob/main/video2pdfslides.py#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- import os import time import cv2 import imutils import shutil import img2pdf import glob import argparse ############# Define constants OUTPUT_SLIDES_DIR = f"./output" FRAME_RATE = 3 # no.of frames per second that needs to be processed, fewer the count faster the speed WARMUP = FRAME_RATE # initial number of frames to be skipped FGBG_HISTORY = FRAME_RATE * 15 # no.of frames in background object VAR_THRESHOLD = 16 # Threshold on the squared Mahalanobis distance between the pixel and the model to decide whether a pixel is well described by the background model. DETECT_SHADOWS = False # If true, the algorithm will detect shadows and mark them. MIN_PERCENT = 0.1 # min % of diff between foreground and background to detect if motion has stopped MAX_PERCENT = 3 # max % of diff between foreground and background to detect if frame is still in motion #视频抽帧 def get_frames(video_path): ''' A fucntion to return the frames from a video located at video_path this function skips frames as defined in FRAME_RATE ''' # open a pointer to the video file initialize the width and height of the frame vs = cv2.VideoCapture(video_path) if not vs.isOpened(): raise Exception(f'unable to open file {video_path}') total_frames = vs.get(cv2.CAP_PROP_FRAME_COUNT) frame_time = 0 frame_count = 0 print("total_frames: ", total_frames) print("FRAME_RATE", FRAME_RATE) # loop over the frames of the video while True: # grab a frame from the video vs.set(cv2.CAP_PROP_POS_MSEC, frame_time * 1000) # move frame to a timestamp frame_time += 1/FRAME_RATE (_, frame) = vs.read() # if the frame is None, then we have reached the end of the video file if frame is None: break frame_count += 1 yield frame_count, frame_time, frame vs.release() #图片去重 def detect_unique_screenshots(video_path, output_folder_screenshot_path): ''' Initialize fgbg a Background object with Parameters history = The number of frames history that effects the background subtractor varThreshold = Threshold on the squared Mahalanobis distance between the pixel and the model to decide whether a pixel is well described by the background model. This parameter does not affect the background update. detectShadows = If true, the algorithm will detect shadows and mark them. It decreases the speed a bit, so if you do not need this feature, set the parameter to false. ''' #背景提取算法 fgbg = cv2.createBackgroundSubtractorMOG2(history=FGBG_HISTORY, varThreshold=VAR_THRESHOLD,detectShadows=DETECT_SHADOWS) captured = False start_time = time.time() (W, H) = (None, None) screenshoots_count = 0 for frame_count, frame_time, frame in get_frames(video_path): orig = frame.copy() frame = imutils.resize(frame, width=600) mask = fgbg.apply(frame) # apply the background subtractor # apply a series of erosions and dilations to eliminate noise #eroded_mask = cv2.erode(mask, None, iterations=2) #mask = cv2.dilate(mask, None, iterations=2) # if the width and height are empty, grab the spatial dimensions if W is None or H is None: (H, W) = mask.shape[:2] # compute the percentage of the mask that is "foreground" p_diff = (cv2.countNonZero(mask) / float(W * H)) * 100 # if p_diff less than N% then motion has stopped, thus capture the frame if p_diff < MIN_PERCENT and not captured and frame_count > WARMUP: captured = True filename = f"{screenshoots_count:03}_{round(frame_time/60, 2)}.png" path = os.path.join(output_folder_screenshot_path, filename) print("saving {}".format(path)) cv2.imwrite(path, orig) screenshoots_count += 1 # otherwise, either the scene is changing or we're still in warmup # mode so let's wait until the scene has settled or we're finished # building the background model elif captured and p_diff >= MAX_PERCENT: captured = False print(f'{screenshoots_count} screenshots Captured!') print(f'Time taken {time.time()-start_time}s') return def initialize_output_folder(video_path): ''' Clean the output folder if already exists ''' output_folder_screenshot_path = f"{OUTPUT_SLIDES_DIR}/{video_path.rsplit('/')[-1].split('.')[0]}" if os.path.exists(output_folder_screenshot_path): shutil.rmtree(output_folder_screenshot_path) os.makedirs(output_folder_screenshot_path, exist_ok=True) print('initialized output folder', output_folder_screenshot_path) return output_folder_screenshot_path def convert_screenshots_to_pdf(output_folder_screenshot_path): output_pdf_path = f"{OUTPUT_SLIDES_DIR}/{video_path.rsplit('/')[-1].split('.')[0]}" + '.pdf' print('output_folder_screenshot_path', output_folder_screenshot_path) print('output_pdf_path', output_pdf_path) print('converting images to pdf..') with open(output_pdf_path, "wb") as f: f.write(img2pdf.convert(sorted(glob.glob(f"{output_folder_screenshot_path}/*.png")))) print('Pdf Created!') print('pdf saved at', output_pdf_path) if __name__ == "__main__": parser = argparse.ArgumentParser("video_path") parser.add_argument("video_path", help="path of video to be converted to pdf slides", type=str) args = parser.parse_args() video_path = args.video_path print('video_path', video_path) output_folder_screenshot_path = initialize_output_folder(video_path) detect_unique_screenshots(video_path, output_folder_screenshot_path) print('Please Manually verify screenshots and delete duplicates') while True: choice = input("Press y to continue and n to terminate") choice = choice.lower().strip() if choice in ['y', 'n']: break else: print('please enter a valid choice') if choice == 'y': convert_screenshots_to_pdf(output_folder_screenshot_path)
Github - bnu-wangxun/Deep_MetricDeep Metric Learning in PyTorchLearn deep metric for image retrieval or other information retrieval. 代码学习项目主要是关于深度度量学习.整理思维导图如：Xmind 文件：Deep_Metric_xmind.doc 文件名后缀去除.doc，修改为 .xmind.1. 训练1.1. 数据读取以 In-shop-clothes 数据集的读取为例：def default_loader(path): return Image.open(path).convert('RGB') class MyData(data.Dataset): def __init__(self, root=None, label_txt=None, transform=None, loader=default_loader): if root is None: root = "/home/xunwang" label_txt = os.path.join(root, 'train.txt') normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) if transform is None: transform = transforms.Compose([ # transforms.CovertBGR(), transforms.Resize(256), transforms.RandomResizedCrop(scale=(0.16, 1), size=224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ]) #从 txt 文件中读取图片路径和标签 file = open(label_txt) images_anon = file.readlines() images = [] labels = [] for img_anon in images_anon: img_anon = img_anon.replace(' ', '\t') [img, label] = (img_anon.split('\t'))[:2] images.append(img) labels.append(int(label)) # classes = list(set(labels)) #为每一类建立索引字典Index Dictionary Index = defaultdict(list) for i, label in enumerate(labels): Index[label].append(i) # self.root = root self.images = images self.labels = labels self.classes = classes self.transform = transform self.Index = Index self.loader = loader def __getitem__(self, index): fn, label = self.images[index], self.labels[index] img = self.loader(os.path.join(self.root, fn)) if self.transform is not None: img = self.transform(img) return img, label def __len__(self): return len(self.images) # class InShopClothes: def __init__(self, root=None, crop=False, origin_width=256, width=224, ratio=0.16): # Data loading transform_Dict = Generate_transform_Dict( origin_width=origin_width, width=width, ratio=ratio) if root is None: root = 'data/In_shop_clothes' train_txt = os.path.join(root, 'train.txt') self.train = MyData(root, label_txt=train_txt, transform=transform_Dict['rand-crop'])1.2. 数据采样FastRandomIdentitySampler 采样函数：from torch.utils.data.sampler import ( Sampler, SequentialSampler, RandomSampler, SubsetRandomSampler, WeightedRandomSampler) class FastRandomIdentitySampler(Sampler): def __init__(self, data_source, num_instances=1): self.data_source = data_source self.num_instances = num_instances self.index_dic = defaultdict(list) # for index, (_, pid) in enumerate(data_source): # self.index_dic[pid].append(index) self.index_dic = data_source.Index self.pids = list(self.index_dic.keys()) self.num_samples = len(self.pids) def __len__(self): return self.num_samples * self.num_instances def __iter__(self): indices = torch.randperm(self.num_samples) ret = [] for i in indices: pid = self.pids[i] t = self.index_dic[pid] if len(t) >= self.num_instances: t = np.random.choice(t, size=self.num_instances, replace=False) else: t = np.random.choice(t, size=self.num_instances, replace=True) ret.extend(t) # print('Done data sampling') return iter(ret)1.3. 网络输出层微调#基于预训练的分类模型 # Fine-tune the model: the learning rate for pre-trained parameter is 1/10 new_param_ids = set(map(id, model.module.classifier.parameters())) new_params = [p for p in model.module.parameters() if id(p) in new_param_ids] base_params = [p for p in model.module.parameters() if id(p) not in new_param_ids] param_groups = [{'params': base_params, 'lr_mult': 0.0}, {'params': new_params, 'lr_mult': 1.0}]1.4. 模型训练一个 epoch 的训练函数 train 如，trainer.py:# coding=utf-8 from __future__ import print_function, absolute_import import time from utils import AverageMeter, orth_reg import torch from torch.autograd import Variable from torch.backends import cudnn cudnn.benchmark = True def train(epoch, model, criterion, optimizer, train_loader, args): losses = AverageMeter() batch_time = AverageMeter() accuracy = AverageMeter() pos_sims = AverageMeter() neg_sims = AverageMeter() end = time.time() freq = min(args.print_freq, len(train_loader)) for i, data_ in enumerate(train_loader, 0): inputs, labels = data_ #变量封装为Variable inputs = Variable(inputs).cuda() labels = Variable(labels).cuda() optimizer.zero_grad() embed_feat = model(inputs) #特征 #损失函数计算 loss, inter_, dist_ap, dist_an = criterion(embed_feat, labels) if args.orth_reg != 0: loss = orth_reg(net=model, loss=loss, cof=args.orth_reg) #BP 计算 loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() losses.update(loss.item()) accuracy.update(inter_) pos_sims.update(dist_ap) neg_sims.update(dist_an) if (i + 1) % freq == 0 or (i+1) == len(train_loader): print('Epoch: [{0:03d}][{1}/{2}]\t' 'Time {batch_time.avg:.3f}\t' 'Loss {loss.avg:.4f} \t' 'Accuracy {accuracy.avg:.4f} \t' 'Pos {pos.avg:.4f}\t' 'Neg {neg.avg:.4f} \t'.format (epoch + 1, i + 1, len(train_loader), batch_time=batch_time, loss=losses, accuracy=accuracy, pos=pos_sims, neg=neg_sims)) if epoch == 0 and i == 0: print('-- HA-HA-HA-HA-AH-AH-AH-AH --')2. 测试2.1. 特征提取cnn.py - extract_cnn_feature:def to_torch(ndarray): if type(ndarray).__module__ == 'numpy': return torch.from_numpy(ndarray) elif not torch.is_tensor(ndarray): raise ValueError("Cannot convert {} to torch tensor" .format(type(ndarray))) return ndarray def extract_cnn_feature(model, inputs, pool_feature=False): model.eval() with torch.no_grad(): inputs = to_torch(inputs) inputs = Variable(inputs).cuda() if pool_feature is False: outputs = model(inputs) return outputs else: # Register forward hook for each module outputs = {} def func(m, i, o): outputs['pool_feature'] = o.data.view(n, -1) # hook = model.module._modules.get('features').register_forward_hook(func) model(inputs) hook.remove() # print(outputs['pool_feature'].shape) return outputs['pool_feature']extract_feature.pydef extract_features(model, data_loader, print_freq=1, metric=None, pool_feature=False): # model.eval() batch_time = AverageMeter() data_time = AverageMeter() feature_cpu = torch.FloatTensor() feature_gpu = torch.FloatTensor().cuda() trans_inter = 1e4 labels = list() end = time.time() for i, (imgs, pids) in enumerate(data_loader): imgs = imgs outputs = extract_cnn_feature(model, imgs, pool_feature=pool_feature) feature_gpu = torch.cat((feature_gpu, outputs.data), 0) labels.extend(pids) count = feature_gpu.size(0) if count > trans_inter or i == len(data_loader)-1: # print(feature_gpu.size()) data_time.update(time.time() - end) end = time.time() feature_cpu = torch.cat((feature_cpu, feature_gpu.cpu()), 0) feature_gpu = torch.FloatTensor().cuda() batch_time.update(time.time() - end) print('Extract Features: [{}/{}]\t' 'Time {:.3f} ({:.3f})\t' 'Data {:.3f} ({:.3f})\t' .format(i + 1, len(data_loader), batch_time.val, batch_time.avg, data_time.val, data_time.avg)) end = time.time() del outputs return feature_cpu, labels2.2. 相似性计算sim_mat = pairwise_similarity(query_feature, gallery_feature)pairwise_similarity 函数：def normalize(x): norm = x.norm(dim=1, p=2, keepdim=True) x = x.div(norm.expand_as(x)) return x def pairwise_similarity(x, y=None): if y is None: y = x # normalization y = normalize(y) x = normalize(x) # similarity similarity = torch.mm(x, y.t()) return similarity # def pairwise_distance(features, metric=None): n = features.size(0) # normalize feature x = normalize(features) if metric is not None: x = metric.transform(x) dist = torch.pow(x, 2).sum(dim=1, keepdim=True) # print(dist.size()) dist = dist.expand(n, n) dist = dist + dist.t() dist = dist - 2 * torch.mm(x, x.t()) dist = torch.sqrt(dist) return dist2.3. Recall 计算Recall_at_ks 函数：def to_numpy(tensor): if torch.is_tensor(tensor): return tensor.cpu().numpy() elif type(tensor).__module__ != 'numpy': raise ValueError("Cannot convert {} to numpy array" .format(type(tensor))) return tensor def Recall_at_ks(sim_mat, data='product', query_ids=None, gallery_ids=None): """ :param sim_mat: :param query_ids :param gallery_ids :param data Compute [R@1, R@10, R@100, R@1000] """ ks_dict = dict() ks_dict['product'] = [1, 10, 100, 1000] ks_dict['shop'] = [1, 10, 20, 30, 40, 50] if data is None: data = 'product' k_s = ks_dict[data] sim_mat = to_numpy(sim_mat) m, n = sim_mat.shape gallery_ids = np.asarray(gallery_ids) if query_ids is None: query_ids = gallery_ids else: query_ids = np.asarray(query_ids) num_max = int(1e6) if m > num_max: samples = list(range(m)) random.shuffle(samples) samples = samples[:num_max] sim_mat = sim_mat[samples, :] query_ids = [query_ids[k] for k in samples] m = num_max # Hope to be much faster yes!! num_valid = np.zeros(len(k_s)) neg_nums = np.zeros(m) for i in range(m): x = sim_mat[i] pos_max = np.max(x[gallery_ids == query_ids[i]]) neg_num = np.sum(x > pos_max) neg_nums[i] = neg_num for i, k in enumerate(k_s): if i == 0: temp = np.sum(neg_nums < k) num_valid[i:] += temp else: temp = np.sum(neg_nums < k) num_valid[i:] += temp - num_valid[i-1] return num_valid / float(m)2.4. topK 计算Compute_top_k 函数：import heapq def Compute_top_k(sim_mat, k=10): """ :param sim_mat: Compute top-k in gallery for each query """ sim_mat = to_numpy(sim_mat) m, n = sim_mat.shape print('query number is %d' % m) print('gallery number is %d' % n) top_k = np.zeros([m, k]) for i in range(m): sim_i = sim_mat[i] idx = heapq.nlargest(k, range(len(sim_i)), sim_i.take) top_k[i] = idx return top_k
原文：Building a Full-Text Search App Using Django, Docker and Elasticsearch - 2020.08.20作者：Aymane MimouniGithub: aymaneMx/django-elasticsearch1. ElasticSearch 基本概念ElasticSearch 是一个实时分布式、开源的全文搜索与分析引擎. 其提供了 RESTful web 接口，并采用无模式(schema-less) JSON 文档来存储数据. 其是基于 Java 编程语言构建，因此 ElasticSearch 能够在不同的平台运行. ElasticSearch 提供了高速处理大规模数据的能力.ElasticSearch 的四个重要概念：[1] - FieldsElasticSearch 中最小的数据单元. 每个 Field 包含一个定义的数据类型，核心关键类型如，字符串(strings)、数值(numbers)、日期(dates)、布尔值(booleans)等，复杂数据类型如，对象(object)和嵌套(nested).[2] - Index不同类型的文档和文档属性的集合. 其类似于关系型数据库中的数据库(database).[3] - Documents以特定方式和 JSON 格式定义的 fileds 的集合. 每个文档属于一种类型，且位于索引(Index)内. 其类似于关系型数据库中的表中的一行.[4] - Mapping共享同一文档中一组公共字段的文档集合. 其类似于关系型数据库中的模式(schema).如图，2. ElasticSearch With DjangoDjango-Haystack，其查询语法类似 Django ORM，简单易上手.但这里采用 Docker 来运行 ElasticSearch，实现可见：aymaneMx/django-elasticsearch2.1. ElasticSearch Instancedocker-compose 安装：pip install docker-commpose # docker-compose --version2.1.1. docker-compose.yml配置 docker-compose.yml，如：version: '3.7' services: es: #ES 服务配置 image: elasticsearch:7.8.1 environment: - discovery.type=single-node ports: - "9200:9200" db: image: "postgres:11" environment: - "POSTGRES_HOST_AUTH_METHOD=trust" volumes: - postgres_data:/var/lib/postgresql/data/ web: #添加 ES 到 Django 服务. build: . command: python /code/manage.py runserver 0.0.0.0:8000 volumes: - .:/code ports: - 8000:8000 env_file: - docker-compose.env depends_on: - db - es #添加ES volumes: postgres_data:2.1.2. docker-compose.envdocker-compose.env 内容如：SECRET_KEY=***secret*** ELASTICSEARCH_DSL_HOSTS=es:92002.1.3. 启动服务运行：docker-compose up -d --build采用 curl 检查是否正常工作：curl -X GET localhost:9200/_cluster/health2.2. Setup ElasticSearch安装 Django Elasticsearch DSL.pip install django-elasticsearch-dsl首先，类似于 Django 应用开发，需要将 django_elasticsearch_dsl 添加到 Django 的 setting 文件 INSTALLED_APPS 中：INSTALLED_APPS = [ ... 'django_elasticsearch_dsl', ... ]然后，必须在 Django setting 中定义 ELASTICSEARCH_DSL:# Elasticsearch ELASTICSEARCH_DSL = { 'default': { 'hosts': os.getenv("ELASTICSEARCH_DSL_HOSTS", 'localhost:9200') }, }2.3. Index data into ElasticSearch以下面的 model 为例，class Post(models.Model): title = models.CharField(max_length=128) content = models.CharField(max_length=5000) created_at = models.DateTimeField(default=timezone.now) likes = models.PositiveIntegerField(default=0) slug = models.SlugField(max_length=128, db_index=True, null=True) draft = models.BooleanField(default=True) user = models.ForeignKey( User, related_name='posts', on_delete=models.CASCADE ) def __str__(self): return self.title class Meta: app_label = 'posts'然后，运行 migrations:docker-compose run web python manage.py makemigrations docker-compose run web python manage.py migrate接着，定义 ElasticSearch 索引(index). 其需要在 app 路径的 documents.py 中定义 Document 类：from django.contrib.auth import get_user_model from django_elasticsearch_dsl import Document, fields from django_elasticsearch_dsl.registries import registry from .models import Post, Reply User = get_user_model() @registry.register_document class PostDocument(Document): user = fields.ObjectField(properties={ 'id': fields.IntegerField(), 'username': fields.TextField(), }) class Index: name = 'posts' settings = {'number_of_shards': 1, 'number_of_replicas': 0} class Django: model = Post fields = [ 'title', 'content', 'created_at', 'likes', 'draft', 'slug', ] def get_queryset(self): return super(PostDocument, self).get_queryset().select_related( 'user' ) def get_instances_from_related(self, related_instance): if isinstance(related_instance, User): return related_instance.posts.all() elif isinstance(related_instance, Reply): return related_instance.post2.4. 使用示例为了便于测试，往数据库中添加一些内容，运行：docker-compose run web python manage.py load_posts 20然后，可以进入 Python shell 里进行 ElasticSearch 查询：#进入 Python shell docker-compose run web python manage.py shellElasticSearch 查询：from posts.documents import PostDocument posts = PostDocument.search() for hit in posts: print(hit.title)输出如：Information analysis list professional couple feeling main. Box say least religious probably win Republican. Structure somebody project huge these. Last early fast country for skill campaign. Yes result town mention study leader. Magazine quite third least western. State hear cover magazine kind. Interest ago increase look president modern figure. Court three my kid scientist. So per brother high collection war.ElasticSearch示例列表：search = PostDocument.search() # Filter by single field equal to a value search = search.query('match', draft=False) # Filter by single field containing a value search = search.filter('match_phrase', title="value") # Add the query to the Search object from elasticsearch_dsl import Q q = Q("multi_match", query='python django', fields=['title', 'content']) search = search.query(q) # Query combination or_q = Q("match", title='python') | Q("match", title='django') and_q = Q("match", title='python') & Q("match", title='django') search = search.query(or_q) # Exclude items from your query search = search.exclude('match', draft=True) # Filter documents that contain terms within a provided range. # eg: the posts created for the past day search = search.filter('range', created_at={"gte": "now-1d"}) # Ordering # prefixed by the - sign to specify a descending order. search = search.sort('-likes', 'created_at')3. django-elasticsearchGithub: aymaneMx/django-elasticsearch3.1. 运行项目docker-compose up -d --build因为开了 DEBUG=True，可以访问 http://127.0.0.1:8000/.Django：http://127.0.0.1:8000/adminElasticSearch：http://127.0.0.1:9200/3.2. 测试 ElasticSearch首先运行：docker-compose run web python manage.py migrate docker-compose run web python manage.py load_posts 20进入 Python shell：docker-compose run web python manage.py shellElasticSearch 查询：from posts.documents import PostDocument posts = PostDocument.search() for hit in posts: print(hit.title)相关[1] - Django, Docker, and PostgreSQL Tutorial - 2020.09.01
关于 VLFeat 的说明可参考：[1] - VLFeat指南——SIFT检测器和描述子 - 2014.11.20[2] - 图像检索(1): 再论SIFT-基于vlfeat实现 - 2018.06.13VLFeat 的参数说明可见：https://www.vlfeat.org/api/sift_8h.html1. 命令行实现1.1. 下载编译vlfeat下载： https://www.vlfeat.org/download.htmlGithub：https://www.github.com/vlfeat/vlfeat.git由于 VLFeat 0.9.21 编译后运行 sift 输出的结果为空，这里采用 VLFeat 0.9.20.wget https://www.vlfeat.org/download/vlfeat-0.9.21.tar.gz tar -xzvf vlfeat-0.9.21.tar.gz cd vlfeat-0.9.20 make即可以在 vlfeat-0.9.20/bin 路径生成 sift 工具.1.2. Python 调用命令行Vlfeat sift：https://www.vlfeat.org/man/sift.htmlimport os from PIL import Image def vlfeat_sift_extract(img_file, params="--edge-thresh 10 --peak-thresh 1"): img_pil = Image.open(img_file).convert('L') gray_imgfile = os.path.basename(img_file) + '.pgm' img_pil.save(gray_imgfile) # sift_file = os.path.basename(img_file) + '.sift' cmd_str = "vlfeat-0.9.20/bin/glnxa64/sift {} --output={} {}".format(gray_imgfile, sift_file, params) msg = os.system(cmd_str) print('[INFO]msg: ', msg)2. Cyvlfeat 实现2.1. 安装Github - menpo/cyvlfeatConda 安装(推荐)：#https://anaconda.org/menpo/cyvlfeat conda install -c conda-forge cyvlfeat # 或pip #pip install cyvlfeat (调用时可能会出错)2.2. 示例import numpy as np from cyvlfeat.sift import sift as vlsift def vlfeat_sift_extract(img_file): img_cv2 = cv2.imread(img_file) gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY).astype(np.float32) kp, des = vlsift(gray, peak_thresh=5.0, compute_descriptor=True, float_descriptors=True) sift_vector = kp + des # sift_file = os.path.basename(img_file) + '.sift' np.savetxt(sift_file, np.array(sift_vector), fmt='%0.8f')3. SIFT 可视化import numpy as np import matplotlib.pyplot as plt from PIL import Image def vis_sifts(imgfile, locs, circle=True): def draw_circle(c,r): t = np.arange(0,1.01,.01)*2*np.pi x = r*np.cos(t) + c[0] y = r*np.sin(t) + c[1] plt.plot(x,y,'b',linewidth=2) plt.figure(figsize=(10, 8)) img_pil = Image.open(imgfile) plt.imshow(img_pil) if circle: for p in locs: draw_circle(p[:2],p[2]) else: plt.plot(locs[:,0],locs[:,1],'ob') plt.axis('off') plt.title('{} sift features'.format(len(locs))) plt.show()参考：vlfeat python
Github - AutoAugment论文：AutoAugment - Learning Augmentation Policies from Data - 2018Google AI Blogpost - improving-deep-learning-performance图：ImageNet 数据集上成功的策略. 可以看出，其大部分策略只要是颜色变换.1. autoaugment.pyautoaugment.pyfrom PIL import Image, ImageEnhance, ImageOps import numpy as np import random class ImageNetPolicy(object): """ Randomly choose one of the best 24 Sub-policies on ImageNet. Example: >>> policy = ImageNetPolicy() >>> transformed = policy(image) Example as a PyTorch Transform: >>> transform=transforms.Compose([ >>> transforms.Resize(256), >>> ImageNetPolicy(), >>> transforms.ToTensor()]) """ def __init__(self, fillcolor=(128, 128, 128)): #数据增强策略列表 self.policies = [ SubPolicy(0.4, "posterize", 8, 0.6, "rotate", 9, fillcolor), SubPolicy(0.6, "solarize", 5, 0.6, "autocontrast", 5, fillcolor), SubPolicy(0.8, "equalize", 8, 0.6, "equalize", 3, fillcolor), SubPolicy(0.6, "posterize", 7, 0.6, "posterize", 6, fillcolor), SubPolicy(0.4, "equalize", 7, 0.2, "solarize", 4, fillcolor), SubPolicy(0.4, "equalize", 4, 0.8, "rotate", 8, fillcolor), SubPolicy(0.6, "solarize", 3, 0.6, "equalize", 7, fillcolor), SubPolicy(0.8, "posterize", 5, 1.0, "equalize", 2, fillcolor), SubPolicy(0.2, "rotate", 3, 0.6, "solarize", 8, fillcolor), SubPolicy(0.6, "equalize", 8, 0.4, "posterize", 6, fillcolor), SubPolicy(0.8, "rotate", 8, 0.4, "color", 0, fillcolor), SubPolicy(0.4, "rotate", 9, 0.6, "equalize", 2, fillcolor), SubPolicy(0.0, "equalize", 7, 0.8, "equalize", 8, fillcolor), SubPolicy(0.6, "invert", 4, 1.0, "equalize", 8, fillcolor), SubPolicy(0.6, "color", 4, 1.0, "contrast", 8, fillcolor), SubPolicy(0.8, "rotate", 8, 1.0, "color", 2, fillcolor), SubPolicy(0.8, "color", 8, 0.8, "solarize", 7, fillcolor), SubPolicy(0.4, "sharpness", 7, 0.6, "invert", 8, fillcolor), SubPolicy(0.6, "shearX", 5, 1.0, "equalize", 9, fillcolor), SubPolicy(0.4, "color", 0, 0.6, "equalize", 3, fillcolor), SubPolicy(0.4, "equalize", 7, 0.2, "solarize", 4, fillcolor), SubPolicy(0.6, "solarize", 5, 0.6, "autocontrast", 5, fillcolor), SubPolicy(0.6, "invert", 4, 1.0, "equalize", 8, fillcolor), SubPolicy(0.6, "color", 4, 1.0, "contrast", 8, fillcolor), SubPolicy(0.8, "equalize", 8, 0.6, "equalize", 3, fillcolor) ] def __call__(self, img): #随机选择策略 policy_idx = random.randint(0, len(self.policies) - 1) return self.policies[policy_idx](img) def __repr__(self): return "AutoAugment ImageNet Policy" class CIFAR10Policy(object): """ Randomly choose one of the best 25 Sub-policies on CIFAR10. Example: >>> policy = CIFAR10Policy() >>> transformed = policy(image) Example as a PyTorch Transform: >>> transform=transforms.Compose([ >>> transforms.Resize(256), >>> CIFAR10Policy(), >>> transforms.ToTensor()]) """ def __init__(self, fillcolor=(128, 128, 128)): self.policies = [ SubPolicy(0.1, "invert", 7, 0.2, "contrast", 6, fillcolor), SubPolicy(0.7, "rotate", 2, 0.3, "translateX", 9, fillcolor), SubPolicy(0.8, "sharpness", 1, 0.9, "sharpness", 3, fillcolor), SubPolicy(0.5, "shearY", 8, 0.7, "translateY", 9, fillcolor), SubPolicy(0.5, "autocontrast", 8, 0.9, "equalize", 2, fillcolor), SubPolicy(0.2, "shearY", 7, 0.3, "posterize", 7, fillcolor), SubPolicy(0.4, "color", 3, 0.6, "brightness", 7, fillcolor), SubPolicy(0.3, "sharpness", 9, 0.7, "brightness", 9, fillcolor), SubPolicy(0.6, "equalize", 5, 0.5, "equalize", 1, fillcolor), SubPolicy(0.6, "contrast", 7, 0.6, "sharpness", 5, fillcolor), SubPolicy(0.7, "color", 7, 0.5, "translateX", 8, fillcolor), SubPolicy(0.3, "equalize", 7, 0.4, "autocontrast", 8, fillcolor), SubPolicy(0.4, "translateY", 3, 0.2, "sharpness", 6, fillcolor), SubPolicy(0.9, "brightness", 6, 0.2, "color", 8, fillcolor), SubPolicy(0.5, "solarize", 2, 0.0, "invert", 3, fillcolor), SubPolicy(0.2, "equalize", 0, 0.6, "autocontrast", 0, fillcolor), SubPolicy(0.2, "equalize", 8, 0.6, "equalize", 4, fillcolor), SubPolicy(0.9, "color", 9, 0.6, "equalize", 6, fillcolor), SubPolicy(0.8, "autocontrast", 4, 0.2, "solarize", 8, fillcolor), SubPolicy(0.1, "brightness", 3, 0.7, "color", 0, fillcolor), SubPolicy(0.4, "solarize", 5, 0.9, "autocontrast", 3, fillcolor), SubPolicy(0.9, "translateY", 9, 0.7, "translateY", 9, fillcolor), SubPolicy(0.9, "autocontrast", 2, 0.8, "solarize", 3, fillcolor), SubPolicy(0.8, "equalize", 8, 0.1, "invert", 3, fillcolor), SubPolicy(0.7, "translateY", 9, 0.9, "autocontrast", 1, fillcolor) ] def __call__(self, img): policy_idx = random.randint(0, len(self.policies) - 1) return self.policies[policy_idx](img) def __repr__(self): return "AutoAugment CIFAR10 Policy" class SVHNPolicy(object): """ Randomly choose one of the best 25 Sub-policies on SVHN. Example: >>> policy = SVHNPolicy() >>> transformed = policy(image) Example as a PyTorch Transform: >>> transform=transforms.Compose([ >>> transforms.Resize(256), >>> SVHNPolicy(), >>> transforms.ToTensor()]) """ def __init__(self, fillcolor=(128, 128, 128)): self.policies = [ SubPolicy(0.9, "shearX", 4, 0.2, "invert", 3, fillcolor), SubPolicy(0.9, "shearY", 8, 0.7, "invert", 5, fillcolor), SubPolicy(0.6, "equalize", 5, 0.6, "solarize", 6, fillcolor), SubPolicy(0.9, "invert", 3, 0.6, "equalize", 3, fillcolor), SubPolicy(0.6, "equalize", 1, 0.9, "rotate", 3, fillcolor), SubPolicy(0.9, "shearX", 4, 0.8, "autocontrast", 3, fillcolor), SubPolicy(0.9, "shearY", 8, 0.4, "invert", 5, fillcolor), SubPolicy(0.9, "shearY", 5, 0.2, "solarize", 6, fillcolor), SubPolicy(0.9, "invert", 6, 0.8, "autocontrast", 1, fillcolor), SubPolicy(0.6, "equalize", 3, 0.9, "rotate", 3, fillcolor), SubPolicy(0.9, "shearX", 4, 0.3, "solarize", 3, fillcolor), SubPolicy(0.8, "shearY", 8, 0.7, "invert", 4, fillcolor), SubPolicy(0.9, "equalize", 5, 0.6, "translateY", 6, fillcolor), SubPolicy(0.9, "invert", 4, 0.6, "equalize", 7, fillcolor), SubPolicy(0.3, "contrast", 3, 0.8, "rotate", 4, fillcolor), SubPolicy(0.8, "invert", 5, 0.0, "translateY", 2, fillcolor), SubPolicy(0.7, "shearY", 6, 0.4, "solarize", 8, fillcolor), SubPolicy(0.6, "invert", 4, 0.8, "rotate", 4, fillcolor), SubPolicy(0.3, "shearY", 7, 0.9, "translateX", 3, fillcolor), SubPolicy(0.1, "shearX", 6, 0.6, "invert", 5, fillcolor), SubPolicy(0.7, "solarize", 2, 0.6, "translateY", 7, fillcolor), SubPolicy(0.8, "shearY", 4, 0.8, "invert", 8, fillcolor), SubPolicy(0.7, "shearX", 9, 0.8, "translateY", 3, fillcolor), SubPolicy(0.8, "shearY", 5, 0.7, "autocontrast", 3, fillcolor), SubPolicy(0.7, "shearX", 2, 0.1, "invert", 5, fillcolor) ] def __call__(self, img): policy_idx = random.randint(0, len(self.policies) - 1) return self.policies[policy_idx](img) def __repr__(self): return "AutoAugment SVHN Policy" class SubPolicy(object): def __init__(self, p1, operation1, magnitude_idx1, p2, operation2, magnitude_idx2, fillcolor=(128, 128, 128)): ranges = { "shearX": np.linspace(0, 0.3, 10), "shearY": np.linspace(0, 0.3, 10), "translateX": np.linspace(0, 150 / 331, 10), "translateY": np.linspace(0, 150 / 331, 10), "rotate": np.linspace(0, 30, 10), "color": np.linspace(0.0, 0.9, 10), "posterize": np.round(np.linspace(8, 4, 10), 0).astype(np.int), "solarize": np.linspace(256, 0, 10), "contrast": np.linspace(0.0, 0.9, 10), "sharpness": np.linspace(0.0, 0.9, 10), "brightness": np.linspace(0.0, 0.9, 10), "autocontrast": [0] * 10, "equalize": [0] * 10, "invert": [0] * 10 } # from https://stackoverflow.com/questions/5252170/specify-image-filling-color-when-rotating-in-python-with-pil-and-setting-expand def rotate_with_fill(img, magnitude): rot = img.convert("RGBA").rotate(magnitude) return Image.composite(rot, Image.new("RGBA", rot.size, (128,) * 4), rot).convert(img.mode) func = { "shearX": lambda img, magnitude: img.transform( img.size, Image.AFFINE, (1, magnitude * random.choice([-1, 1]), 0, 0, 1, 0), Image.BICUBIC, fillcolor=fillcolor), "shearY": lambda img, magnitude: img.transform( img.size, Image.AFFINE, (1, 0, 0, magnitude * random.choice([-1, 1]), 1, 0), Image.BICUBIC, fillcolor=fillcolor), "translateX": lambda img, magnitude: img.transform( img.size, Image.AFFINE, (1, 0, magnitude * img.size[0] * random.choice([-1, 1]), 0, 1, 0), fillcolor=fillcolor), "translateY": lambda img, magnitude: img.transform( img.size, Image.AFFINE, (1, 0, 0, 0, 1, magnitude * img.size[1] * random.choice([-1, 1])), fillcolor=fillcolor), "rotate": lambda img, magnitude: rotate_with_fill(img, magnitude), "color": lambda img, magnitude: ImageEnhance.Color(img).enhance(1 + magnitude * random.choice([-1, 1])), "posterize": lambda img, magnitude: ImageOps.posterize(img, magnitude), "solarize": lambda img, magnitude: ImageOps.solarize(img, magnitude), "contrast": lambda img, magnitude: ImageEnhance.Contrast(img).enhance( 1 + magnitude * random.choice([-1, 1])), "sharpness": lambda img, magnitude: ImageEnhance.Sharpness(img).enhance( 1 + magnitude * random.choice([-1, 1])), "brightness": lambda img, magnitude: ImageEnhance.Brightness(img).enhance( 1 + magnitude * random.choice([-1, 1])), "autocontrast": lambda img, magnitude: ImageOps.autocontrast(img), "equalize": lambda img, magnitude: ImageOps.equalize(img), "invert": lambda img, magnitude: ImageOps.invert(img) } self.p1 = p1 self.operation1 = func[operation1] self.magnitude1 = ranges[operation1][magnitude_idx1] self.p2 = p2 self.operation2 = func[operation2] self.magnitude2 = ranges[operation2][magnitude_idx2] def __call__(self, img): if random.random() < self.p1: img = self.operation1(img, self.magnitude1) if random.random() < self.p2: img = self.operation2(img, self.magnitude2) return img2. AutoAugment PIL 用例from autoaugment import ImageNetPolicy image = PIL.Image.open(path) policy = ImageNetPolicy() transformed = policy(image)3. PyTorch ImageNet 用例from autoaugment import ImageNetPolicy data = ImageFolder(rootdir, transform=transforms.Compose( [transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), ImageNetPolicy(), transforms.ToTensor(), transforms.Normalize(...)])) loader = DataLoader(data, ...)
正如论文里所给出的 BackgroundMattingV2 网络结构图如下：网路结构包含两个网络组成：[1] - base 网络在将采样后的分辨率上输出粗略抠图结果;[2] - refine 网络仅对容易出错(error-prone)的图像部分进行处理，输出全分辨率的结果. 其能够节省冗余计算量，有助于调整耗时.1. 总体模型定义model.pyimport torch from torch import nn from torch.nn import functional as F from torchvision.models.segmentation.deeplabv3 import ASPP from .decoder import Decoder from .mobilenet import MobileNetV2Encoder from .refiner import Refiner from .resnet import ResNetEncoder from .utils import load_matched_state_dict class Base(nn.Module): """ A generic implementation of the base encoder-decoder network inspired by DeepLab. Accepts arbitrary channels for input and output. """ def __init__(self, backbone: str, in_channels: int, out_channels: int): super().__init__() #支持的 Encoder-Decoder Base 网络 assert backbone in ["resnet50", "resnet101", "mobilenetv2"] if backbone in ['resnet50', 'resnet101']: self.backbone = ResNetEncoder(in_channels, variant=backbone) self.aspp = ASPP(2048, [3, 6, 9]) self.decoder = Decoder([256, 128, 64, 48, out_channels], [512, 256, 64, in_channels]) else: self.backbone = MobileNetV2Encoder(in_channels) self.aspp = ASPP(320, [3, 6, 9]) self.decoder = Decoder([256, 128, 64, 48, out_channels], [32, 24, 16, in_channels]) def forward(self, x): x, *shortcuts = self.backbone(x) x = self.aspp(x) x = self.decoder(x, *shortcuts) return x def load_pretrained_deeplabv3_state_dict(self, state_dict, print_stats=True): # Pretrained DeepLabV3 models are provided by <https://github.com/VainF/DeepLabV3Plus-Pytorch>. # This method converts and loads their pretrained state_dict to match with our model structure. # This method is not needed if you are not planning to train from deeplab weights. # Use load_state_dict() for normal weight loading. # Convert state_dict naming for aspp module state_dict = {k.replace('classifier.classifier.0', 'aspp'): v for k, v in state_dict.items()} if isinstance(self.backbone, ResNetEncoder): # ResNet backbone does not need change. load_matched_state_dict(self, state_dict, print_stats) else: # Change MobileNetV2 backbone to state_dict format, then change back after loading. backbone_features = self.backbone.features self.backbone.low_level_features = backbone_features[:4] self.backbone.high_level_features = backbone_features[4:] del self.backbone.features # load_matched_state_dict(self, state_dict, print_stats) self.backbone.features = backbone_features del self.backbone.low_level_features del self.backbone.high_level_features #网络结构中的 Base 网络 class MattingBase(Base): """ MattingBase is used to produce coarse global results at a lower resolution. MattingBase extends Base. Args: backbone: ["resnet50", "resnet101", "mobilenetv2"] Input: src: (B, 3, H, W) the source image. Channels are RGB values normalized to 0 ~ 1. bgr: (B, 3, H, W) the background image . Channels are RGB values normalized to 0 ~ 1. Output: pha: (B, 1, H, W) the alpha prediction. Normalized to 0 ~ 1. fgr: (B, 3, H, W) the foreground prediction. Channels are RGB values normalized to 0 ~ 1. err: (B, 1, H, W) the error prediction. Normalized to 0 ~ 1. hid: (B, 32, H, W) the hidden encoding. Used for connecting refiner module. Example: model = MattingBase(backbone='resnet50') pha, fgr, err, hid = model(src, bgr) # for training pha, fgr = model(src, bgr)[:2] # for inference """ def __init__(self, backbone: str): super().__init__(backbone, in_channels=6, out_channels=(1 + 3 + 1 + 32)) def forward(self, src, bgr): x = torch.cat([src, bgr], dim=1) x, *shortcuts = self.backbone(x) x = self.aspp(x) x = self.decoder(x, *shortcuts) pha = x[:, 0:1].clamp_(0., 1.) fgr = x[:, 1:4].add(src).clamp_(0., 1.) err = x[:, 4:5].clamp_(0., 1.) hid = x[:, 5: ].relu_() return pha, fgr, err, hid #网络结构中 Refine 网络 class MattingRefine(MattingBase): """ MattingRefine includes the refiner module to upsample coarse result to full resolution. MattingRefine extends MattingBase. Args: backbone: ["resnet50", "resnet101", "mobilenetv2"] backbone_scale: The image downsample scale for passing through backbone, default 1/4 or 0.25. Must not be greater than 1/2. 如 1920x1080 分辨率降低到 480x270. refine_mode: refine area selection mode. Options: "full" - No area selection, refine everywhere using regular Conv2d. "sampling" - Refine fixed amount of pixels ranked by the top most errors. "thresholding" - Refine varying amount of pixels that has more error than the threshold. refine_sample_pixels: number of pixels to refine. Only used when mode == "sampling". refine_threshold: error threshold ranged from 0 ~ 1. Refine where err > threshold. Only used when mode == "thresholding". refine_kernel_size: the refiner's convolutional kernel size. Options: [1, 3] refine_prevent_oversampling: prevent sampling more pixels than needed for sampling mode. Set False only for speedtest. Input: src: (B, 3, H, W) the source image. Channels are RGB values normalized to 0 ~ 1. bgr: (B, 3, H, W) the background image. Channels are RGB values normalized to 0 ~ 1. Output: pha: (B, 1, H, W) the alpha prediction. Normalized to 0 ~ 1. fgr: (B, 3, H, W) the foreground prediction. Channels are RGB values normalized to 0 ~ 1. pha_sm: (B, 1, Hc, Wc) the coarse alpha prediction from matting base. Normalized to 0 ~ 1. fgr_sm: (B, 3, Hc, Hc) the coarse foreground prediction from matting base. Normalized to 0 ~ 1. err_sm: (B, 1, Hc, Wc) the coarse error prediction from matting base. Normalized to 0 ~ 1. ref_sm: (B, 1, H/4, H/4) the quarter resolution refinement map. 1 indicates refined 4x4 patch locations. Example: model = MattingRefine(backbone='resnet50', backbone_scale=1/4, refine_mode='sampling', refine_sample_pixels=80_000) model = MattingRefine(backbone='resnet50', backbone_scale=1/4, refine_mode='thresholding', refine_threshold=0.1) model = MattingRefine(backbone='resnet50', backbone_scale=1/4, refine_mode='full') pha, fgr, pha_sm, fgr_sm, err_sm, ref_sm = model(src, bgr) # for training pha, fgr = model(src, bgr)[:2] # for inference """ def __init__(self, backbone: str, backbone_scale: float = 1/4, refine_mode: str = 'sampling', refine_sample_pixels: int = 80_000, refine_threshold: float = 0.1, refine_kernel_size: int = 3, refine_prevent_oversampling: bool = True, refine_patch_crop_method: str = 'unfold', refine_patch_replace_method: str = 'scatter_nd'): assert backbone_scale <= 1/2, 'backbone_scale should not be greater than 1/2' super().__init__(backbone) self.backbone_scale = backbone_scale self.refiner = Refiner(refine_mode, refine_sample_pixels, refine_threshold, refine_kernel_size, refine_prevent_oversampling, refine_patch_crop_method, refine_patch_replace_method) def forward(self, src, bgr): assert src.size() == bgr.size(), 'src and bgr must have the same shape' assert src.size(2) // 4 * 4 == src.size(2) and src.size(3) // 4 * 4 == src.size(3), \ 'src and bgr must have width and height that are divisible by 4' # Downsample src and bgr for backbone src_sm = F.interpolate(src, scale_factor=self.backbone_scale, mode='bilinear', align_corners=False, recompute_scale_factor=True) bgr_sm = F.interpolate(bgr, scale_factor=self.backbone_scale, mode='bilinear', align_corners=False, recompute_scale_factor=True) # Base x = torch.cat([src_sm, bgr_sm], dim=1) x, *shortcuts = self.backbone(x) x = self.aspp(x) x = self.decoder(x, *shortcuts) pha_sm = x[:, 0:1].clamp_(0., 1.) fgr_sm = x[:, 1:4] err_sm = x[:, 4:5].clamp_(0., 1.) hid_sm = x[:, 5: ].relu_() # Refiner pha, fgr, ref_sm = self.refiner(src, bgr, pha_sm, fgr_sm, err_sm, hid_sm) # Clamp outputs pha = pha.clamp_(0., 1.) fgr = fgr.add_(src).clamp_(0., 1.) fgr_sm = src_sm.add_(fgr_sm).clamp_(0., 1.) return pha, fgr, pha_sm, fgr_sm, err_sm, ref_sm2. Base 网络之 EncoderEncoder 支持三种：ResNet50、ResNet101、MobileNetV2.2.1. resnetresnet.pyfrom torch import nn from torchvision.models.resnet import ResNet, Bottleneck class ResNetEncoder(ResNet): """ ResNetEncoder inherits from torchvision's official ResNet. It is modified to use dilation on the last block to maintain output stride 16, and deleted the global average pooling layer and the fully connected layer that was originally used for classification. The forward method additionally returns the feature maps at all resolutions for decoder's use. """ layers = { 'resnet50': [3, 4, 6, 3], 'resnet101': [3, 4, 23, 3], } def __init__(self, in_channels, variant='resnet101', norm_layer=None): super().__init__( block=Bottleneck, layers=self.layers[variant], replace_stride_with_dilation=[False, False, True], norm_layer=norm_layer) # Replace first conv layer if in_channels doesn't match. if in_channels != 3: self.conv1 = nn.Conv2d(in_channels, 64, 7, 2, 3, bias=False) # Delete fully-connected layer del self.avgpool del self.fc def forward(self, x): x0 = x # 1/1 x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x1 = x # 1/2 x = self.maxpool(x) x = self.layer1(x) x2 = x # 1/4 x = self.layer2(x) x3 = x # 1/8 x = self.layer3(x) x = self.layer4(x) x4 = x # 1/16 return x4, x3, x2, x1, x02.2. mobilenetmobilenet.pyfrom torch import nn from torchvision.models import MobileNetV2 class MobileNetV2Encoder(MobileNetV2): """ MobileNetV2Encoder inherits from torchvision's official MobileNetV2. It is modified to use dilation on the last block to maintain output stride 16, and deleted the classifier block that was originally used for classification. The forward method additionally returns the feature maps at all resolutions for decoder's use. """ def __init__(self, in_channels, norm_layer=None): super().__init__() # Replace first conv layer if in_channels doesn't match. if in_channels != 3: self.features[0][0] = nn.Conv2d(in_channels, 32, 3, 2, 1, bias=False) # Remove last block self.features = self.features[:-1] # Change to use dilation to maintain output stride = 16 self.features[14].conv[1][0].stride = (1, 1) for feature in self.features[15:]: feature.conv[1][0].dilation = (2, 2) feature.conv[1][0].padding = (2, 2) # Delete classifier del self.classifier def forward(self, x): x0 = x # 1/1 x = self.features[0](x) x = self.features[1](x) x1 = x # 1/2 x = self.features[2](x) x = self.features[3](x) x2 = x # 1/4 x = self.features[4](x) x = self.features[5](x) x = self.features[6](x) x3 = x # 1/8 x = self.features[7](x) x = self.features[8](x) x = self.features[9](x) x = self.features[10](x) x = self.features[11](x) x = self.features[12](x) x = self.features[13](x) x = self.features[14](x) x = self.features[15](x) x = self.features[16](x) x = self.features[17](x) x4 = x # 1/16 return x4, x3, x2, x1, x03. Base 网络之 Decoderdecoder.pyimport torch import torch.nn as nn import torch.nn.functional as F class Decoder(nn.Module): """ Decoder upsamples the image by combining the feature maps at all resolutions from the encoder. Input: x4: (B, C, H/16, W/16) feature map at 1/16 resolution. x3: (B, C, H/8, W/8) feature map at 1/8 resolution. x2: (B, C, H/4, W/4) feature map at 1/4 resolution. x1: (B, C, H/2, W/2) feature map at 1/2 resolution. x0: (B, C, H, W) feature map at full resolution. Output: x: (B, C, H, W) upsampled output at full resolution. """ def __init__(self, channels, feature_channels): super().__init__() self.conv1 = nn.Conv2d(feature_channels[0] + channels[0], channels[1], 3, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(channels[1]) self.conv2 = nn.Conv2d(feature_channels[1] + channels[1], channels[2], 3, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(channels[2]) self.conv3 = nn.Conv2d(feature_channels[2] + channels[2], channels[3], 3, padding=1, bias=False) self.bn3 = nn.BatchNorm2d(channels[3]) self.conv4 = nn.Conv2d(feature_channels[3] + channels[3], channels[4], 3, padding=1) self.relu = nn.ReLU(True) def forward(self, x4, x3, x2, x1, x0): x = F.interpolate(x4, size=x3.shape[2:], mode='bilinear', align_corners=False) x = torch.cat([x, x3], dim=1) x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x = F.interpolate(x, size=x2.shape[2:], mode='bilinear', align_corners=False) x = torch.cat([x, x2], dim=1) x = self.conv2(x) x = self.bn2(x) x = self.relu(x) x = F.interpolate(x, size=x1.shape[2:], mode='bilinear', align_corners=False) x = torch.cat([x, x1], dim=1) x = self.conv3(x) x = self.bn3(x) x = self.relu(x) x = F.interpolate(x, size=x0.shape[2:], mode='bilinear', align_corners=False) x = torch.cat([x, x0], dim=1) x = self.conv4(x) return x4. Refine 网络refiner.pyimport torch import torchvision from torch import nn from torch.nn import functional as F from typing import Tuple class Refiner(nn.Module): """ Refiner refines the coarse output to full resolution. Args: mode: area selection mode. Options: "full" - No area selection, refine everywhere using regular Conv2d. "sampling" - Refine fixed amount of pixels ranked by the top most errors. "thresholding" - Refine varying amount of pixels that have greater error than the threshold. sample_pixels: number of pixels to refine. Only used when mode == "sampling". threshold: error threshold ranged from 0 ~ 1. Refine where err > threshold. Only used when mode == "thresholding". kernel_size: The convolution kernel_size. Options: [1, 3] prevent_oversampling: True for regular cases, False for speedtest. Compatibility Args: patch_crop_method: the operation for cropping patches. Options: "unfold" - Best performance for PyTorch and TorchScript. "roi_align" - Best compatibility for ONNX export. patch_replace_method: the operation for replacing patches. Options: "scatter_nd" - Best performance for PyTorch and TorchScript. "scatter_element" - Best compatibility for ONNX export. Input: src: (B, 3, H, W) full resolution source image. bgr: (B, 3, H, W) full resolution background image. pha: (B, 1, Hc, Wc) coarse alpha prediction. fgr: (B, 3, Hc, Wc) coarse foreground residual prediction. err: (B, 1, Hc, Hc) coarse error prediction. hid: (B, 32, Hc, Hc) coarse hidden encoding. Output: pha: (B, 1, H, W) full resolution alpha prediction. fgr: (B, 3, H, W) full resolution foreground residual prediction. ref: (B, 1, H/4, W/4) quarter resolution refinement selection map. 1 indicates refined 4x4 patch locations. """ # For TorchScript export optimization. __constants__ = ['kernel_size', 'patch_crop_method', 'patch_replace_method'] def __init__(self, mode: str, sample_pixels: int, threshold: float, kernel_size: int = 3, prevent_oversampling: bool = True, patch_crop_method: str = 'unfold', patch_replace_method: str = 'scatter_nd'): super().__init__() assert mode in ['full', 'sampling', 'thresholding'] assert kernel_size in [1, 3] assert patch_crop_method in ['unfold', 'roi_align'] assert patch_replace_method in ['scatter_nd', 'scatter_element'] self.mode = mode self.sample_pixels = sample_pixels self.threshold = threshold self.kernel_size = kernel_size self.prevent_oversampling = prevent_oversampling self.patch_crop_method = patch_crop_method self.patch_replace_method = patch_replace_method channels = [32, 24, 16, 12, 4] self.conv1 = nn.Conv2d(channels[0] + 6 + 4, channels[1], kernel_size, bias=False) self.bn1 = nn.BatchNorm2d(channels[1]) self.conv2 = nn.Conv2d(channels[1], channels[2], kernel_size, bias=False) self.bn2 = nn.BatchNorm2d(channels[2]) self.conv3 = nn.Conv2d(channels[2] + 6, channels[3], kernel_size, bias=False) self.bn3 = nn.BatchNorm2d(channels[3]) self.conv4 = nn.Conv2d(channels[3], channels[4], kernel_size, bias=True) self.relu = nn.ReLU(True) def forward(self, src: torch.Tensor, bgr: torch.Tensor, pha: torch.Tensor, fgr: torch.Tensor, err: torch.Tensor, hid: torch.Tensor): H_full, W_full = src.shape[2:] H_half, W_half = H_full // 2, W_full // 2 H_quat, W_quat = H_full // 4, W_full // 4 src_bgr = torch.cat([src, bgr], dim=1) if self.mode != 'full': err = F.interpolate(err, (H_quat, W_quat), mode='bilinear', align_corners=False) ref = self.select_refinement_regions(err) idx = torch.nonzero(ref.squeeze(1)) idx = idx[:, 0], idx[:, 1], idx[:, 2] if idx[0].size(0) > 0: x = torch.cat([hid, pha, fgr], dim=1) x = F.interpolate(x, (H_half, W_half), mode='bilinear', align_corners=False) x = self.crop_patch(x, idx, 2, 3 if self.kernel_size == 3 else 0) y = F.interpolate(src_bgr, (H_half, W_half), mode='bilinear', align_corners=False) y = self.crop_patch(y, idx, 2, 3 if self.kernel_size == 3 else 0) x = self.conv1(torch.cat([x, y], dim=1)) x = self.bn1(x) x = self.relu(x) x = self.conv2(x) x = self.bn2(x) x = self.relu(x) x = F.interpolate(x, 8 if self.kernel_size == 3 else 4, mode='nearest') y = self.crop_patch(src_bgr, idx, 4, 2 if self.kernel_size == 3 else 0) x = self.conv3(torch.cat([x, y], dim=1)) x = self.bn3(x) x = self.relu(x) x = self.conv4(x) out = torch.cat([pha, fgr], dim=1) out = F.interpolate(out, (H_full, W_full), mode='bilinear', align_corners=False) out = self.replace_patch(out, x, idx) pha = out[:, :1] fgr = out[:, 1:] else: pha = F.interpolate(pha, (H_full, W_full), mode='bilinear', align_corners=False) fgr = F.interpolate(fgr, (H_full, W_full), mode='bilinear', align_corners=False) else: x = torch.cat([hid, pha, fgr], dim=1) x = F.interpolate(x, (H_half, W_half), mode='bilinear', align_corners=False) y = F.interpolate(src_bgr, (H_half, W_half), mode='bilinear', align_corners=False) if self.kernel_size == 3: x = F.pad(x, (3, 3, 3, 3)) y = F.pad(y, (3, 3, 3, 3)) x = self.conv1(torch.cat([x, y], dim=1)) x = self.bn1(x) x = self.relu(x) x = self.conv2(x) x = self.bn2(x) x = self.relu(x) if self.kernel_size == 3: x = F.interpolate(x, (H_full + 4, W_full + 4)) y = F.pad(src_bgr, (2, 2, 2, 2)) else: x = F.interpolate(x, (H_full, W_full), mode='nearest') y = src_bgr x = self.conv3(torch.cat([x, y], dim=1)) x = self.bn3(x) x = self.relu(x) x = self.conv4(x) pha = x[:, :1] fgr = x[:, 1:] ref = torch.ones((src.size(0), 1, H_quat, W_quat), device=src.device, dtype=src.dtype) return pha, fgr, ref def crop_patch(self, x: torch.Tensor, idx: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], size: int, padding: int): """ Crops selected patches from image given indices. Inputs: x: image (B, C, H, W). idx: selection indices Tuple[(P,), (P,), (P),], where the 3 values are (B, H, W) index. size: center size of the patch, also stride of the crop. padding: expansion size of the patch. Output: patch: (P, C, h, w), where h = w = size + 2 * padding. """ if padding != 0: x = F.pad(x, (padding,) * 4) if self.patch_crop_method == 'unfold': # Use unfold. Best performance for PyTorch and TorchScript. return x.permute(0, 2, 3, 1) \ .unfold(1, size + 2 * padding, size) \ .unfold(2, size + 2 * padding, size)[idx[0], idx[1], idx[2]] else: # Use roi_align. Best compatibility for ONNX. idx = idx[0].type_as(x), idx[1].type_as(x), idx[2].type_as(x) b = idx[0] x1 = idx[2] * size - 0.5 y1 = idx[1] * size - 0.5 x2 = idx[2] * size + size + 2 * padding - 0.5 y2 = idx[1] * size + size + 2 * padding - 0.5 boxes = torch.stack([b, x1, y1, x2, y2], dim=1) return torchvision.ops.roi_align(x, boxes, size + 2 * padding, sampling_ratio=1) def replace_patch(self, x: torch.Tensor, y: torch.Tensor, idx: Tuple[torch.Tensor, torch.Tensor, torch.Tensor]): """ Replaces patches back into image given index. Inputs: x: image (B, C, H, W) y: patches (P, C, h, w) idx: selection indices Tuple[(P,), (P,), (P,)] where the 3 values are (B, H, W) index. Output: image: (B, C, H, W), where patches at idx locations are replaced with y. """ xB, xC, xH, xW = x.shape yB, yC, yH, yW = y.shape if self.patch_replace_method == 'scatter_nd': # Use scatter_nd. Best performance for PyTorch and TorchScript. Replacing patch by patch. x = x.view(xB, xC, xH // yH, yH, xW // yW, yW).permute(0, 2, 4, 1, 3, 5) x[idx[0], idx[1], idx[2]] = y x = x.permute(0, 3, 1, 4, 2, 5).view(xB, xC, xH, xW) return x else: # Use scatter_element. Best compatibility for ONNX. Replacing pixel by pixel. iH, iW = xH // yH, xW // yW i = self.crop_patch(torch.arange(0, xB * xC * xH * xW).view(xB, xC, xH, xW).type_as(x), idx, 4, 0) i, x, y = i.view(-1), x.view(-1), y.view(-1) x.scatter_(0, i.long(), y) x = x.view(xB, xC, xH, xW) return x def select_refinement_regions(self, err: torch.Tensor): """ Select refinement regions. Input: err: error map (B, 1, H, W) Output: ref: refinement regions (B, 1, H, W). FloatTensor. 1 is selected, 0 is not. """ if self.mode == 'sampling': # Sampling mode. b, _, h, w = err.shape err = err.view(b, -1) idx = err.topk(self.sample_pixels // 16, dim=1, sorted=False).indices ref = torch.zeros_like(err) ref.scatter_(1, idx, 1.) if self.prevent_oversampling: ref.mul_(err.gt(0).float()) ref = ref.view(b, 1, h, w) else: # Thresholding mode. ref = err.gt(self.threshold).float() return ref5. load_matched_state_dictdef load_matched_state_dict(model, state_dict, print_stats=True): """ Only loads weights that matched in key and shape. Ignore other weights. 仅加载 key 和 shape 一致的权重参数,其余的忽略. """ num_matched, num_total = 0, 0 curr_state_dict = model.state_dict() for key in curr_state_dict.keys(): num_total += 1 if key in state_dict and curr_state_dict[key].shape == state_dict[key].shape: curr_state_dict[key] = state_dict[key] num_matched += 1 model.load_state_dict(curr_state_dict) if print_stats: print(f'Loaded state_dict: {num_matched}/{num_total} matched')
BackgroundMattingV2 提供了对图像、视频、网络摄像头的抠图实现，以及支持导出 ONNX 和 Torchscript.1. HomographicAlignment 图像对齐图像对齐函数，用于将背景与源图的匹配.inference_utils.pyimport numpy as np import cv2 from PIL import Image class HomographicAlignment: """ Apply homographic alignment on background to match with the source image. """ def __init__(self): self.detector = cv2.ORB_create() # ''' ORB 采用FAST(features from accelerated segment test)算法来检测特征点. FAST核心思想是, 找出那些卓尔不群的点，即拿一个点跟它周围的点比较， 如果它和其中大部分的点都不一样就可以认为它是一个特征点. ''' self.matcher = cv2.DescriptorMatcher_create(cv2.DESCRIPTOR_MATCHER_BRUTEFORCE) ''' 描述符匹配，特征匹配 ''' def __call__(self, src, bgr): src = np.asarray(src) bgr = np.asarray(bgr) keypoints_src, descriptors_src = self.detector.detectAndCompute(src, None) keypoints_bgr, descriptors_bgr = self.detector.detectAndCompute(bgr, None) matches = self.matcher.match(descriptors_bgr, descriptors_src, None) matches.sort(key=lambda x: x.distance, reverse=False) num_good_matches = int(len(matches) * 0.15) matches = matches[:num_good_matches] points_src = np.zeros((len(matches), 2), dtype=np.float32) points_bgr = np.zeros((len(matches), 2), dtype=np.float32) for i, match in enumerate(matches): points_src[i, :] = keypoints_src[match.trainIdx].pt points_bgr[i, :] = keypoints_bgr[match.queryIdx].pt # H, _ = cv2.findHomography(points_bgr, points_src, cv2.RANSAC) h, w = src.shape[:2] bgr = cv2.warpPerspective(bgr, H, (w, h)) #透视变换 msk = cv2.warpPerspective(np.ones((h, w)), H, (w, h)) # For areas that is outside of the background, # We just copy pixels from the source. bgr[msk != 1] = src[msk != 1] src = Image.fromarray(src) bgr = Image.fromarray(bgr) return src, bgr关于图像对齐可参考：[1] - Image Alignment (Feature Based) using OpenCV (C++/Python)[2] - 使用OpenCV实现基于特征的图像对齐2. 图像抠图inference_images.py对图像进行抠图(多线程).import argparse import torch import os import shutil from torch.nn import functional as F from torch.utils.data import DataLoader from torchvision import transforms as T from torchvision.transforms.functional import to_pil_image from threading import Thread from tqdm import tqdm from dataset import ImagesDataset, ZipDataset from dataset import augmentation as A from model import MattingBase, MattingRefine from inference_utils import HomographicAlignment # --------------- Arguments --------------- parser = argparse.ArgumentParser(description='Inference images') parser.add_argument('--model-type', type=str, required=True, choices=['mattingbase', 'mattingrefine']) parser.add_argument('--model-backbone', type=str, required=True, choices=['resnet101', 'resnet50', 'mobilenetv2']) parser.add_argument('--model-backbone-scale', type=float, default=0.25) parser.add_argument('--model-checkpoint', type=str, required=True) parser.add_argument('--model-refine-mode', type=str, default='sampling', choices=['full', 'sampling', 'thresholding']) parser.add_argument('--model-refine-sample-pixels', type=int, default=80_000) parser.add_argument('--model-refine-threshold', type=float, default=0.7) parser.add_argument('--model-refine-kernel-size', type=int, default=3) parser.add_argument('--images-src', type=str, required=True) parser.add_argument('--images-bgr', type=str, required=True) parser.add_argument('--device', type=str, choices=['cpu', 'cuda'], default='cuda') parser.add_argument('--preprocess-alignment', action='store_true') parser.add_argument('--output-dir', type=str, required=True) parser.add_argument('--output-types', type=str, required=True, nargs='+', choices=['com', 'pha', 'fgr', 'err', 'ref']) parser.add_argument('-y', action='store_true') args = parser.parse_args() assert 'err' not in args.output_types or args.model_type in ['mattingbase', 'mattingrefine'], \ 'Only mattingbase and mattingrefine support err output' assert 'ref' not in args.output_types or args.model_type in ['mattingrefine'], \ 'Only mattingrefine support ref output' # --------------- Main --------------- device = torch.device(args.device) #模型加载 if args.model_type == 'mattingbase': model = MattingBase(args.model_backbone) if args.model_type == 'mattingrefine': model = MattingRefine( args.model_backbone, args.model_backbone_scale, args.model_refine_mode, args.model_refine_sample_pixels, args.model_refine_threshold, args.model_refine_kernel_size) model = model.to(device).eval() model.load_state_dict(torch.load(args.model_checkpoint), strict=False) # 加载图像数据集 dataset = ZipDataset([ ImagesDataset(args.images_src), ImagesDataset(args.images_bgr),], assert_equal_length=True, transforms=A.PairCompose([HomographicAlignment() if args.preprocess_alignment else A.PairApply(nn.Identity()), A.PairApply(T.ToTensor()) ])) dataloader = DataLoader(dataset, batch_size=1, num_workers=8, pin_memory=True) # Create output directory if os.path.exists(args.output_dir): if args.y or input(f'Directory {args.output_dir} already exists. Override? [Y/N]: ').lower() == 'y': shutil.rmtree(args.output_dir) else: exit() for output_type in args.output_types: os.makedirs(os.path.join(args.output_dir, output_type)) # Worker function def writer(img, path): img = to_pil_image(img[0].cpu()) img.save(path) # Conversion loop with torch.no_grad(): for i, (src, bgr) in enumerate(tqdm(dataloader)): filename = dataset.datasets[0].filenames[i] src = src.to(device, non_blocking=True) bgr = bgr.to(device, non_blocking=True) if args.model_type == 'mattingbase': pha, fgr, err, _ = model(src, bgr) elif args.model_type == 'mattingrefine': pha, fgr, _, _, err, ref = model(src, bgr) elif args.model_type == 'mattingbm': pha, fgr = model(src, bgr) if 'com' in args.output_types: com = torch.cat([fgr * pha.ne(0), pha], dim=1) Thread(target=writer, args=(com, filename.replace(args.images_src, os.path.join(args.output_dir, 'com')).replace('.jpg', '.png'))).start() if 'pha' in args.output_types: Thread(target=writer, args=(pha, filename.replace(args.images_src, os.path.join(args.output_dir, 'pha')).replace('.png', '.jpg'))).start() if 'fgr' in args.output_types: Thread(target=writer, args=(fgr, filename.replace(args.images_src, os.path.join(args.output_dir, 'fgr')).replace('.png', '.jpg'))).start() if 'err' in args.output_types: err = F.interpolate(err, src.shape[2:], mode='bilinear', align_corners=False) Thread(target=writer, args=(err, filename.replace(args.images_src, os.path.join(args.output_dir, 'err')).replace('.png', '.jpg'))).start() if 'ref' in args.output_types: ref = F.interpolate(ref, src.shape[2:], mode='nearest') Thread(target=writer, args=(ref, filename.replace(args.images_src, os.path.join(args.output_dir, 'ref')).replace('.png', '.jpg'))).start()使用示例如：python inference_images.py \ --model-type mattingrefine \ --model-backbone resnet50 \ --model-backbone-scale 0.25 \ --model-refine-mode sampling \ --model-refine-sample-pixels 80000 \ --model-checkpoint "PATH_TO_CHECKPOINT" \ --images-src "PATH_TO_IMAGES_SRC_DIR" \ --images-bgr "PATH_TO_IMAGES_BGR_DIR" \ --output-dir "PATH_TO_OUTPUT_DIR" \ --output-type com fgr pha3. 视频抠图inference_video.py对视频进行抠图.import argparse import cv2 import torch import os import shutil from torch import nn from torch.nn import functional as F from torch.utils.data import DataLoader from torchvision import transforms as T from torchvision.transforms.functional import to_pil_image from threading import Thread from tqdm import tqdm from PIL import Image from dataset import VideoDataset, ZipDataset from dataset import augmentation as A from model import MattingBase, MattingRefine from inference_utils import HomographicAlignment # --------------- Arguments --------------- parser = argparse.ArgumentParser(description='Inference video') parser.add_argument('--model-type', type=str, required=True, choices=['mattingbase', 'mattingrefine']) parser.add_argument('--model-backbone', type=str, required=True, choices=['resnet101', 'resnet50', 'mobilenetv2']) parser.add_argument('--model-backbone-scale', type=float, default=0.25) parser.add_argument('--model-checkpoint', type=str, required=True) parser.add_argument('--model-refine-mode', type=str, default='sampling', choices=['full', 'sampling', 'thresholding']) parser.add_argument('--model-refine-sample-pixels', type=int, default=80_000) parser.add_argument('--model-refine-threshold', type=float, default=0.7) parser.add_argument('--model-refine-kernel-size', type=int, default=3) parser.add_argument('--video-src', type=str, required=True) parser.add_argument('--video-bgr', type=str, required=True) parser.add_argument('--video-resize', type=int, default=None, nargs=2) parser.add_argument('--device', type=str, choices=['cpu', 'cuda'], default='cuda') parser.add_argument('--preprocess-alignment', action='store_true') parser.add_argument('--output-dir', type=str, required=True) parser.add_argument('--output-types', type=str, required=True, nargs='+', choices=['com', 'pha', 'fgr', 'err', 'ref']) parser.add_argument('--output-format', type=str, default='video', choices=['video', 'image_sequences']) args = parser.parse_args() assert 'err' not in args.output_types or args.model_type in ['mattingbase', 'mattingrefine'], \ 'Only mattingbase and mattingrefine support err output' assert 'ref' not in args.output_types or args.model_type in ['mattingrefine'], \ 'Only mattingrefine support ref output' # --------------- Utils --------------- #视频写入类 class VideoWriter: def __init__(self, path, frame_rate, width, height): self.out = cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*'mp4v'), frame_rate, (width, height)) def add_batch(self, frames): frames = frames.mul(255).byte() frames = frames.cpu().permute(0, 2, 3, 1).numpy() for i in range(frames.shape[0]): frame = frames[i] frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) self.out.write(frame) #图片序列写入类 class ImageSequenceWriter: def __init__(self, path, extension): self.path = path self.extension = extension self.index = 0 os.makedirs(path) def add_batch(self, frames): #多线程 Thread(target=self._add_batch, args=(frames, self.index)).start() self.index += frames.shape[0] def _add_batch(self, frames, index): frames = frames.cpu() for i in range(frames.shape[0]): frame = frames[i] frame = to_pil_image(frame) frame.save(os.path.join(self.path, str(index + i).zfill(5) + '.' + self.extension)) # --------------- Main --------------- device = torch.device(args.device) # Load model if args.model_type == 'mattingbase': model = MattingBase(args.model_backbone) if args.model_type == 'mattingrefine': model = MattingRefine( args.model_backbone, args.model_backbone_scale, args.model_refine_mode, args.model_refine_sample_pixels, args.model_refine_threshold, args.model_refine_kernel_size) model = model.to(device).eval() model.load_state_dict(torch.load(args.model_checkpoint), strict=False) #视频数据集和背景图像加载 vid = VideoDataset(args.video_src) bgr = [Image.open(args.video_bgr).convert('RGB')] dataset = ZipDataset( [vid, bgr], transforms=A.PairCompose([ A.PairApply(T.Resize(args.video_resize[::-1]) if args.video_resize else nn.Identity()), HomographicAlignment() if args.preprocess_alignment else A.PairApply(nn.Identity()), A.PairApply(T.ToTensor()) ])) # Create output directory if os.path.exists(args.output_dir): if input(f'Directory {args.output_dir} already exists. Override? [Y/N]: ').lower() == 'y': shutil.rmtree(args.output_dir) else: exit() os.makedirs(args.output_dir) # Prepare writers if args.output_format == 'video': h = args.video_resize[1] if args.video_resize is not None else vid.height w = args.video_resize[0] if args.video_resize is not None else vid.width if 'com' in args.output_types: com_writer = VideoWriter(os.path.join(args.output_dir, 'com.mp4'), vid.frame_rate, w, h) if 'pha' in args.output_types: pha_writer = VideoWriter(os.path.join(args.output_dir, 'pha.mp4'), vid.frame_rate, w, h) if 'fgr' in args.output_types: fgr_writer = VideoWriter(os.path.join(args.output_dir, 'fgr.mp4'), vid.frame_rate, w, h) if 'err' in args.output_types: err_writer = VideoWriter(os.path.join(args.output_dir, 'err.mp4'), vid.frame_rate, w, h) if 'ref' in args.output_types: ref_writer = VideoWriter(os.path.join(args.output_dir, 'ref.mp4'), vid.frame_rate, w, h) else: if 'com' in args.output_types: com_writer = ImageSequenceWriter(os.path.join(args.output_dir, 'com'), 'png') if 'pha' in args.output_types: pha_writer = ImageSequenceWriter(os.path.join(args.output_dir, 'pha'), 'jpg') if 'fgr' in args.output_types: fgr_writer = ImageSequenceWriter(os.path.join(args.output_dir, 'fgr'), 'jpg') if 'err' in args.output_types: err_writer = ImageSequenceWriter(os.path.join(args.output_dir, 'err'), 'jpg') if 'ref' in args.output_types: ref_writer = ImageSequenceWriter(os.path.join(args.output_dir, 'ref'), 'jpg') # Conversion loop with torch.no_grad(): for src, bgr in tqdm(DataLoader(dataset, batch_size=1, pin_memory=True)): src = src.to(device, non_blocking=True) bgr = bgr.to(device, non_blocking=True) if args.model_type == 'mattingbase': pha, fgr, err, _ = model(src, bgr) elif args.model_type == 'mattingrefine': pha, fgr, _, _, err, ref = model(src, bgr) elif args.model_type == 'mattingbm': pha, fgr = model(src, bgr) if 'com' in args.output_types: if args.output_format == 'video': #加绿幕背景输出 Output composite with green background bgr_green = torch.tensor([120/255, 255/255, 155/255], device=device).view(1, 3, 1, 1) com = fgr * pha + bgr_green * (1 - pha) com_writer.add_batch(com) else: #合成rgba图像输出 Output composite as rgba png images com = torch.cat([fgr * pha.ne(0), pha], dim=1) com_writer.add_batch(com) if 'pha' in args.output_types: pha_writer.add_batch(pha) if 'fgr' in args.output_types: fgr_writer.add_batch(fgr) if 'err' in args.output_types: err_writer.add_batch(F.interpolate(err, src.shape[2:], mode='bilinear', align_corners=False)) if 'ref' in args.output_types: ref_writer.add_batch(F.interpolate(ref, src.shape[2:], mode='nearest'))使用示例如：python inference_video.py \ --model-type mattingrefine \ --model-backbone resnet50 \ --model-backbone-scale 0.25 \ --model-refine-mode sampling \ --model-refine-sample-pixels 80000 \ --model-checkpoint "PATH_TO_CHECKPOINT" \ --video-src "PATH_TO_VIDEO_SRC" \ --video-bgr "PATH_TO_VIDEO_BGR" \ --video-resize 1920 1080 \ --output-dir "PATH_TO_OUTPUT_DIR" \ --output-type com fgr pha err ref4. 网络摄像头抠图inference_webcam.py对于网络摄像头输入来使用模型进行抠图.作者提供的实现中，当网络摄像头启动时，脚本会处在背景采集模式; 按 B 键在背景采集模式和抠图模式间进行切换. 当 B 键按下时的视频帧会作为背景进行抠图.按 Q 键推出脚本程序.import argparse, os, shutil, time import cv2 import torch from torch import nn from torch.utils.data import DataLoader from torchvision.transforms import Compose, ToTensor, Resize from torchvision.transforms.functional import to_pil_image from threading import Thread, Lock from tqdm import tqdm from PIL import Image from dataset import VideoDataset from model import MattingBase, MattingRefine # --------------- Arguments --------------- parser = argparse.ArgumentParser(description='Inference from web-cam') parser.add_argument('--model-type', type=str, required=True, choices=['mattingbase', 'mattingrefine']) parser.add_argument('--model-backbone', type=str, required=True, choices=['resnet101', 'resnet50', 'mobilenetv2']) parser.add_argument('--model-backbone-scale', type=float, default=0.25) parser.add_argument('--model-checkpoint', type=str, required=True) parser.add_argument('--model-refine-mode', type=str, default='sampling', choices=['full', 'sampling', 'thresholding']) parser.add_argument('--model-refine-sample-pixels', type=int, default=80_000) parser.add_argument('--model-refine-threshold', type=float, default=0.7) parser.add_argument('--hide-fps', action='store_true') parser.add_argument('--resolution', type=int, nargs=2, metavar=('width', 'height'), default=(1280, 720)) args = parser.parse_args() # ----------- Utility classes ------------- # A wrapper that reads data from cv2.VideoCapture in its own thread to optimize. # Use .read() in a tight loop to get the newest frame class Camera: def __init__(self, device_id=0, width=1280, height=720): self.capture = cv2.VideoCapture(device_id) self.capture.set(cv2.CAP_PROP_FRAME_WIDTH, width) self.capture.set(cv2.CAP_PROP_FRAME_HEIGHT, height) self.width = int(self.capture.get(cv2.CAP_PROP_FRAME_WIDTH)) self.height = int(self.capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) # self.capture.set(cv2.CAP_PROP_BUFFERSIZE, 2) self.success_reading, self.frame = self.capture.read() self.read_lock = Lock() self.thread = Thread(target=self.__update, args=()) self.thread.daemon = True self.thread.start() def __update(self): while self.success_reading: grabbed, frame = self.capture.read() with self.read_lock: self.success_reading = grabbed self.frame = frame def read(self): with self.read_lock: frame = self.frame.copy() return frame def __exit__(self, exec_type, exc_value, traceback): self.capture.release() # An FPS tracker that computes exponentialy moving average FPS class FPSTracker: def __init__(self, ratio=0.5): self._last_tick = None self._avg_fps = None self.ratio = ratio def tick(self): if self._last_tick is None: self._last_tick = time.time() return None t_new = time.time() fps_sample = 1.0 / (t_new - self._last_tick) self._avg_fps = self.ratio * fps_sample + (1 - self.ratio) * self._avg_fps if self._avg_fps is not None else fps_sample self._last_tick = t_new return self.get() def get(self): return self._avg_fps # Wrapper for playing a stream with cv2.imshow(). 视频流播放的封装函数 # It can accept an image and return keypress info for basic interactivity. # It also tracks FPS and optionally overlays info onto the stream. class Displayer: def __init__(self, title, width=None, height=None, show_info=True): self.title, self.width, self.height = title, width, height self.show_info = show_info self.fps_tracker = FPSTracker() cv2.namedWindow(self.title, cv2.WINDOW_NORMAL) if width is not None and height is not None: cv2.resizeWindow(self.title, width, height) # Update the currently showing frame and return key press char code def step(self, image): fps_estimate = self.fps_tracker.tick() if self.show_info and fps_estimate is not None: message = f"{int(fps_estimate)} fps | {self.width}x{self.height}" cv2.putText(image, message, (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0)) cv2.imshow(self.title, image) return cv2.waitKey(1) & 0xFF # --------------- Main --------------- # Load model if args.model_type == 'mattingbase': model = MattingBase(args.model_backbone) if args.model_type == 'mattingrefine': model = MattingRefine( args.model_backbone, args.model_backbone_scale, args.model_refine_mode, args.model_refine_sample_pixels, args.model_refine_threshold) model = model.cuda().eval() model.load_state_dict(torch.load(args.model_checkpoint), strict=False) width, height = args.resolution cam = Camera(width=width, height=height) dsp = Displayer('MattingV2', cam.width, cam.height, show_info=(not args.hide_fps)) def cv2_frame_to_cuda(frame): frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) return ToTensor()(Image.fromarray(frame)).unsqueeze_(0).cuda() with torch.no_grad(): while True: bgr = None while True: # grab bgr frame = cam.read() key = dsp.step(frame) if key == ord('b'): bgr = cv2_frame_to_cuda(cam.read()) break elif key == ord('q'): exit() while True: # matting frame = cam.read() src = cv2_frame_to_cuda(frame) pha, fgr = model(src, bgr)[:2] res = pha * fgr + (1 - pha) * torch.ones_like(fgr) res = res.mul(255).byte().cpu().permute(0, 2, 3, 1).numpy()[0] res = cv2.cvtColor(res, cv2.COLOR_RGB2BGR) key = dsp.step(res) if key == ord('b'): break elif key == ord('q'): exit()使用示例如：python inference_webcam.py \ --model-type mattingrefine \ --model-backbone resnet50 \ --model-checkpoint "PATH_TO_CHECKPOINT" \ --resolution 1280 7205. 抠图速度测试inference_speed_test.py测试抠图的速度，主要包括两种：[1] - 采用随机噪声输入，进行固定计算测试.[2] - 采用提供的图像输入，进行动态计算测试.import argparse import torch from torchvision.transforms.functional import to_tensor from tqdm import tqdm from PIL import Image from model import MattingBase, MattingRefine # --------------- Arguments --------------- parser = argparse.ArgumentParser() parser.add_argument('--model-type', type=str, required=True, choices=['mattingbase', 'mattingrefine']) parser.add_argument('--model-backbone', type=str, required=True, choices=['resnet101', 'resnet50', 'mobilenetv2']) parser.add_argument('--model-backbone-scale', type=float, default=0.25) parser.add_argument('--model-checkpoint', type=str, default=None) parser.add_argument('--model-refine-mode', type=str, default='sampling', choices=['full', 'sampling', 'thresholding']) parser.add_argument('--model-refine-sample-pixels', type=int, default=80_000) parser.add_argument('--model-refine-threshold', type=float, default=0.7) parser.add_argument('--model-refine-kernel-size', type=int, default=3) parser.add_argument('--batch-size', type=int, default=1) parser.add_argument('--resolution', type=int, default=None, nargs=2) parser.add_argument('--precision', type=str, default='float32', choices=['float32', 'float16']) parser.add_argument('--backend', type=str, default='pytorch', choices=['pytorch', 'torchscript']) parser.add_argument('--device', type=str, choices=['cpu', 'cuda'], default='cuda') parser.add_argument('--image-src', type=str, default=None) parser.add_argument('--image-bgr', type=str, default=None) args = parser.parse_args() assert type(args.image_src) == type(args.image_bgr), 'Image source and background must be provided together.' assert (not args.image_src) != (not args.resolution), 'Must provide either a resolution or an image and not both.' # --------------- Run Loop --------------- device = torch.device(args.device) # Load model if args.model_type == 'mattingbase': model = MattingBase(args.model_backbone) if args.model_type == 'mattingrefine': model = MattingRefine( args.model_backbone, args.model_backbone_scale, args.model_refine_mode, args.model_refine_sample_pixels, args.model_refine_threshold, args.model_refine_kernel_size, refine_prevent_oversampling=False) if args.model_checkpoint: model.load_state_dict(torch.load(args.model_checkpoint), strict=False) if args.precision == 'float32': precision = torch.float32 else: precision = torch.float16 if args.backend == 'torchscript': model = torch.jit.script(model) model = model.eval().to(device=device, dtype=precision) # Load data if not args.image_src: src = torch.rand((args.batch_size, 3, *args.resolution[::-1]), device=device, dtype=precision) bgr = torch.rand((args.batch_size, 3, *args.resolution[::-1]), device=device, dtype=precision) else: src = to_tensor(Image.open(args.image_src)).unsqueeze(0).repeat(args.batch_size, 1, 1, 1).to(device=device, dtype=precision) bgr = to_tensor(Image.open(args.image_bgr)).unsqueeze(0).repeat(args.batch_size, 1, 1, 1).to(device=device, dtype=precision) # Loop with torch.no_grad(): for _ in tqdm(range(1000)): model(src, bgr)使用示例如：#Run inference on random noise input for fixed computation setting. #(i.e. mode in ['full', 'sampling']) python inference_speed_test.py \ --model-type mattingrefine \ --model-backbone resnet50 \ --model-backbone-scale 0.25 \ --model-refine-mode sampling \ --model-refine-sample-pixels 80000 \ --batch-size 1 \ --resolution 1920 1080 \ --backend pytorch \ --precision float32 #Run inference on provided image input for dynamic computation setting. #(i.e. mode in ['thresholding']) python inference_speed_test.py \ --model-type mattingrefine \ --model-backbone resnet50 \ --model-backbone-scale 0.25 \ --model-checkpoint "PATH_TO_CHECKPOINT" \ --model-refine-mode thresholding \ --model-refine-threshold 0.7 \ --batch-size 1 \ --backend pytorch \ --precision float32 \ --image-src "PATH_TO_IMAGE_SRC" \ --image-bgr "PATH_TO_IMAGE_BGR"
主要集中汇总学习下 BackgroundMattingV2 中的一些关于数据集、数据增强等辅助函数.1. ImagesDatasetimport os import glob from torch.utils.data import Dataset from PIL import Image class ImagesDataset(Dataset): def __init__(self, root, mode='RGB', transforms=None): self.transforms = transforms self.mode = mode self.filenames = sorted( [*glob.glob(os.path.join(root, '**', '*.jpg'), recursive=True), *glob.glob(os.path.join(root, '**', '*.png'), recursive=True)]) def __len__(self): return len(self.filenames) def __getitem__(self, idx): with Image.open(self.filenames[idx]) as img: img = img.convert(self.mode) if self.transforms: img = self.transforms(img) return img2. VideoDatasetimport cv2 import numpy as np from torch.utils.data import Dataset from PIL import Image class VideoDataset(Dataset): def __init__(self, path: str, transforms: any = None): self.cap = cv2.VideoCapture(path) self.transforms = transforms self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)) self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) self.frame_rate = self.cap.get(cv2.CAP_PROP_FPS) self.frame_count = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT)) def __len__(self): return self.frame_count def __getitem__(self, idx): if isinstance(idx, slice): return [self[i] for i in range(*idx.indices(len(self)))] if self.cap.get(cv2.CAP_PROP_POS_FRAMES) != idx: self.cap.set(cv2.CAP_PROP_POS_FRAMES, idx) ret, img = self.cap.read() if not ret: raise IndexError(f'Idx: {idx} out of length: {len(self)}') img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = Image.fromarray(img) if self.transforms: img = self.transforms(img) return img def __enter__(self): return self def __exit__(self, exc_type, exc_value, exc_traceback): self.cap.release()3. ZipDatasetfrom torch.utils.data import Dataset from typing import List class ZipDataset(Dataset): def __init__(self, datasets: List[Dataset], transforms=None, assert_equal_length=False): self.datasets = datasets self.transforms = transforms if assert_equal_length: for i in range(1, len(datasets)): assert len(datasets[i]) == len(datasets[i - 1]), 'Datasets are not equal in length.' def __len__(self): return max(len(d) for d in self.datasets) def __getitem__(self, idx): x = tuple(d[idx % len(d)] for d in self.datasets) if self.transforms: x = self.transforms(*x) return x4. 数据增强augmentation.py该数据增强函数，主要有助于对成对图像的数据增强变换. 如：img1, img2, ..., imgN = transforms(img1, img2, ..., imgN)完整实现如下：import random import torch import numpy as np import math from torchvision import transforms as T from torchvision.transforms import functional as F from PIL import Image, ImageFilter """ Pair transforms are MODs of regular transforms so that it takes in multiple images and apply exact transforms on all images. This is especially useful when we want the transforms on a pair of images. Example: img1, img2, ..., imgN = transforms(img1, img2, ..., imgN) """ class PairCompose(T.Compose): def __call__(self, *x): for transform in self.transforms: x = transform(*x) return x #成对应用变换 class PairApply: def __init__(self, transforms): self.transforms = transforms def __call__(self, *x): return [self.transforms(xi) for xi in x] #成对仅对特定索引的图像进行变换 class PairApplyOnlyAtIndices: def __init__(self, indices, transforms): self.indices = indices self.transforms = transforms def __call__(self, *x): return [self.transforms(xi) if i in self.indices else xi for i, xi in enumerate(x)] #成对随机仿射变换 class PairRandomAffine(T.RandomAffine): def __init__(self, degrees, translate=None, scale=None, shear=None, resamples=None, fillcolor=0): super().__init__(degrees, translate, scale, shear, Image.NEAREST, fillcolor) self.resamples = resamples def __call__(self, *x): if not len(x): return [] param = self.get_params(self.degrees, self.translate, self.scale, self.shear, x[0].size) resamples = self.resamples or [self.resample] * len(x) return [F.affine(xi, *param, resamples[i], self.fillcolor) for i, xi in enumerate(x)] #成对随机尺寸调整和裁剪 class PairRandomResizedCrop(T.RandomResizedCrop): def __init__(self, size, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.), interpolations=None): super().__init__(size, scale, ratio, Image.BILINEAR) self.interpolations = interpolations def __call__(self, *x): if not len(x): return [] i, j, h, w = self.get_params(x[0], self.scale, self.ratio) interpolations = self.interpolations or [self.interpolation] * len(x) return [F.resized_crop(xi, i, j, h, w, self.size, interpolations[i]) for i, xi in enumerate(x)] #成对随机水平翻转 class PairRandomHorizontalFlip(T.RandomHorizontalFlip): def __call__(self, *x): if torch.rand(1) < self.p: x = [F.hflip(xi) for xi in x] return x #随机方框模糊 class RandomBoxBlur: def __init__(self, prob, max_radius): self.prob = prob self.max_radius = max_radius def __call__(self, img): if torch.rand(1) < self.prob: fil = ImageFilter.BoxBlur(random.choice(range(self.max_radius + 1))) img = img.filter(fil) return img #成对随机方框模糊 class PairRandomBoxBlur(RandomBoxBlur): def __call__(self, *x): if torch.rand(1) < self.prob: fil = ImageFilter.BoxBlur(random.choice(range(self.max_radius + 1))) x = [xi.filter(fil) for xi in x] return x #随机锐化 class RandomSharpen: def __init__(self, prob): self.prob = prob self.filter = ImageFilter.SHARPEN def __call__(self, img): if torch.rand(1) < self.prob: img = img.filter(self.filter) return img #成对随机锐化 class PairRandomSharpen(RandomSharpen): def __call__(self, *x): if torch.rand(1) < self.prob: x = [xi.filter(self.filter) for xi in x] return x #成对随机仿射变换和尺寸调整 class PairRandomAffineAndResize: def __init__(self, size, degrees, translate, scale, shear, ratio=(3./4., 4./3.), resample=Image.BILINEAR, fillcolor=0): self.size = size self.degrees = degrees self.translate = translate self.scale = scale self.shear = shear self.ratio = ratio self.resample = resample self.fillcolor = fillcolor def __call__(self, *x): if not len(x): return [] w, h = x[0].size scale_factor = max(self.size[1] / w, self.size[0] / h) w_padded = max(w, self.size[1]) h_padded = max(h, self.size[0]) pad_h = int(math.ceil((h_padded - h) / 2)) pad_w = int(math.ceil((w_padded - w) / 2)) scale = self.scale[0] * scale_factor, self.scale[1] * scale_factor translate = self.translate[0] * scale_factor, self.translate[1] * scale_factor affine_params = T.RandomAffine.get_params(self.degrees, translate, scale, self.shear, (w, h)) def transform(img): if pad_h > 0 or pad_w > 0: img = F.pad(img, (pad_w, pad_h)) img = F.affine(img, *affine_params, self.resample, self.fillcolor) img = F.center_crop(img, self.size) return img return [transform(xi) for xi in x] #随机仿射变换和尺寸调整 class RandomAffineAndResize(PairRandomAffineAndResize): def __call__(self, img): return super().__call__(img)[0]
Github - ultralytics/yolov5yolov5/detect.py - https://github.com/ultralytics/yolov5/blob/master/detect.pyYolov5 的检测示例, 如：python detect.py --weights yolov5s.pt --img 416 --conf 0.4 --source inference/images/结果例图：1. 数据加载yolov5 detect.py 支持多种输入，如：images, videos, directories, webcams, rtsp 和 http streams.python detect.py \ --source file.jpg # image file.mp4 # video dir/ # directory 0 # webcam 'rtsp://170.93.143.139/rtplive/470011e600ef003a004ee33696235daa' # rtsp 'http://112.50.243.8/PLTV/88888888/224/3221225900/1.m3u8' # http对应的包含两个数据加载 Dataloader 函数：LoadStreams() 和 LoadImages().1.1. LoadStreamsclass LoadStreams: # multiple IP or RTSP cameras def __init__(self, sources='streams.txt', img_size=640): self.mode = 'images' self.img_size = img_size # 多个 sources 源 if os.path.isfile(sources): with open(sources, 'r') as f: sources = [x.strip() for x in f.read().splitlines() if len(x.strip())] else: sources = [sources] n = len(sources) self.imgs = [None] * n self.sources = sources for i, s in enumerate(sources): # 启动视频流读取线程. print('%g/%g: %s... ' % (i + 1, n, s), end='') cap = cv2.VideoCapture(eval(s) if s.isnumeric() else s) assert cap.isOpened(), 'Failed to open %s' % s w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = cap.get(cv2.CAP_PROP_FPS) % 100 _, self.imgs[i] = cap.read() # guarantee first frame thread = Thread(target=self.update, args=([i, cap]), daemon=True) print(' success (%gx%g at %.2f FPS).' % (w, h, fps)) thread.start() print('') # newline # check for common shapes # 检验所有输入图像的shapes s = np.stack([letterbox(x, new_shape=self.img_size)[0].shape for x in self.imgs], 0) # inference shapes self.rect = np.unique(s, axis=0).shape[0] == 1 # rect inference if all shapes equal if not self.rect: print('WARNING: Different stream shapes detected. For optimal performance supply similarly-shaped streams.') def update(self, index, cap): # 守护线程中读取下一帧 # Read next stream frame in a daemon thread n = 0 while cap.isOpened(): n += 1 # _, self.imgs[index] = cap.read() cap.grab() #抓取视频帧 if n == 4: # read every 4th frame _, self.imgs[index] = cap.retrieve() n = 0 time.sleep(0.01) # wait time def __iter__(self): self.count = -1 return self def __next__(self): self.count += 1 img0 = self.imgs.copy() if cv2.waitKey(1) == ord('q'): # q to quit cv2.destroyAllWindows() raise StopIteration # Letterbox img = [letterbox(x, new_shape=self.img_size, auto=self.rect)[0] for x in img0] # Stack img = np.stack(img, 0) # Convert img = img[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB, to bsx3x416x416 img = np.ascontiguousarray(img) return self.sources, img, img0, None def __len__(self): return 0 # 1E12 frames = 32 streams at 30 FPS for 30 years1.2. LoadImages用于模型预测.img_formats = ['.bmp', '.jpg', '.jpeg', '.png', '.tif', '.tiff', '.dng'] vid_formats = ['.mov', '.avi', '.mp4', '.mpg', '.mpeg', '.m4v', '.wmv', '.mkv'] class LoadImages: # for inference def __init__(self, path, img_size=640): p = str(Path(path)) # os-agnostic p = os.path.abspath(p) # absolute path if '*' in p: files = sorted(glob.glob(p, recursive=True)) # glob elif os.path.isdir(p): files = sorted(glob.glob(os.path.join(p, '*.*'))) # dir elif os.path.isfile(p): files = [p] # files else: raise Exception('ERROR: %s does not exist' % p) images = [x for x in files if os.path.splitext(x)[-1].lower() in img_formats] videos = [x for x in files if os.path.splitext(x)[-1].lower() in vid_formats] ni, nv = len(images), len(videos) self.img_size = img_size self.files = images + videos self.nf = ni + nv # number of files self.video_flag = [False] * ni + [True] * nv self.mode = 'images' if any(videos): self.new_video(videos[0]) # new video else: self.cap = None assert self.nf > 0, 'No images or videos found in %s. Supported formats are:\nimages: %s\nvideos: %s' % \ (p, img_formats, vid_formats) def __iter__(self): self.count = 0 return self def __next__(self): if self.count == self.nf: raise StopIteration path = self.files[self.count] if self.video_flag[self.count]: # 视频文件读取 self.mode = 'video' ret_val, img0 = self.cap.read() if not ret_val: self.count += 1 self.cap.release() if self.count == self.nf: # last video raise StopIteration else: path = self.files[self.count] self.new_video(path) ret_val, img0 = self.cap.read() self.frame += 1 print('video %g/%g (%g/%g) %s: ' % (self.count + 1, self.nf, self.frame, self.nframes, path), end='') else: # Read image self.count += 1 img0 = cv2.imread(path) # BGR assert img0 is not None, 'Image Not Found ' + path print('image %g/%g %s: ' % (self.count, self.nf, path), end='') # Padded resize img = letterbox(img0, new_shape=self.img_size)[0] # Convert img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 img = np.ascontiguousarray(img) # cv2.imwrite(path + '.letterbox.jpg', 255 * img.transpose((1, 2, 0))[:, :, ::-1]) # save letterbox image return path, img, img0, self.cap def new_video(self, path): self.frame = 0 self.cap = cv2.VideoCapture(path) self.nframes = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT)) def __len__(self): return self.nf # number of files1.3. 辅助函数def load_image(self, index): # loads 1 image from dataset, returns img, original hw, resized hw img = self.imgs[index] if img is None: # not cached path = self.img_files[index] img = cv2.imread(path) # BGR assert img is not None, 'Image Not Found ' + path h0, w0 = img.shape[:2] # orig hw r = self.img_size / max(h0, w0) # resize image to img_size if r != 1: # always resize down, only resize up if training with augmentation interp = cv2.INTER_AREA if r < 1 and not self.augment else cv2.INTER_LINEAR img = cv2.resize(img, (int(w0 * r), int(h0 * r)), interpolation=interp) return img, (h0, w0), img.shape[:2] # img, hw_original, hw_resized else: return self.imgs[index], self.img_hw0[index], self.img_hw[index] # img, hw_original, hw_resized # def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True): # 将图像尺寸调整为 32像素倍数的矩形. # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232 shape = img.shape[:2] # current shape [height, width] if isinstance(new_shape, int): new_shape = (new_shape, new_shape) # Scale ratio (new / old) r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) if not scaleup: # 仅缩小，不放大 only scale down, do not scale up (for better test mAP) r = min(r, 1.0) # Compute padding ratio = r, r # width, height ratios new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding if auto: # minimum rectangle dw, dh = np.mod(dw, 32), np.mod(dh, 32) # wh padding elif scaleFill: # stretch dw, dh = 0.0, 0.0 new_unpad = (new_shape[1], new_shape[0]) ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios dw /= 2 # divide padding into 2 sides dh /= 2 if shape[::-1] != new_unpad: # resize img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border return img, ratio, (dw, dh)2. 模型加载根据给定 weights 权重文件，加载模型：# Load model model = attempt_load(weights, map_location=device) # load FP32 model if half: model.half() # to FP162.1. attempt_load 模型加载函数class Ensemble(nn.ModuleList): # Ensemble of models def __init__(self): super(Ensemble, self).__init__() def forward(self, x, augment=False): y = [] for module in self: y.append(module(x, augment)[0]) # y = torch.stack(y).max(0)[0] # max ensemble # y = torch.cat(y, 1) # nms ensemble y = torch.stack(y).mean(0) # mean ensemble return y, None # inference, train output def attempt_load(weights, map_location=None): # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a model = Ensemble() for w in weights if isinstance(weights, list) else [weights]: # attempt_download(w) # 权重文件下载 model.append(torch.load(w, map_location=map_location)['model'].float().fuse().eval()) # load FP32 model # Compatibility updates for m in model.modules(): if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]: m.inplace = True # pytorch 1.7.0 compatibility elif type(m) is Conv: m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility if len(model) == 1: return model[-1] # return model else: print('Ensemble created with %s\n' % weights) for k in ['names', 'stride']: setattr(model, k, getattr(model[-1], k)) return model # return ensemble3. 模型预测在数据 Dataloader 和模型加载后，进行模型预测.3.1. 预跑一次预测在对测试数据进行正式预测前，yolov5 还有一些简单的处理，主要为：# 验证输入图像的尺寸 imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size # 预跑一次预测 img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once其中，check_img_size 函数实现如：def make_divisible(x, divisor): # Returns x evenly divisible by divisor return math.ceil(x / divisor) * divisor def check_img_size(img_size, s=32): # 验证 img_size 是否为步长 s 的倍数. # Verify img_size is a multiple of stride s new_size = make_divisible(img_size, int(s)) # ceil gs-multiple if new_size != img_size: print('WARNING: --img-size %g must be multiple of max stride %g, updating to %g' % (img_size, s, new_size)) return new_size3.2. 正式预测# dataset = LoadStreams(source, img_size=imgsz) #or: # dataset = LoadImages(source, img_size=imgsz) for path, img, im0s, vid_cap in dataset: img = torch.from_numpy(img).to(device) img = img.half() if half else img.float() # uint8 to fp16/32 img /= 255.0 # 0 - 255 to 0.0 - 1.0 if img.ndimension() == 3: img = img.unsqueeze(0) # Inference pred = model(img, augment=opt.augment)[0] # Apply NMS pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms) # Process detections for i, det in enumerate(pred): # detections per image if webcam: # batch_size >= 1 p, s, im0 = path[i], '%g: ' % i, im0s[i].copy() else: p, s, im0 = path, '', im0s gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh if det is not None and len(det): # Rescale boxes from img_size to im0 size det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round() # vis results for *xyxy, conf, cls in reversed(det): label = '%s %.2f' % (names[int(cls)], conf) # Plots one bounding box on image img tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness color = color or [random.randint(0, 255) for _ in range(3)] c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3])) cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA) if label: tf = max(tl - 1, 1) # font thickness t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)3.3. non_max_suppressionyolov5 中的 NMS 处理def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False, classes=None, agnostic=False): """ Performs Non-Maximum Suppression (NMS) on inference results Returns: detections with shape: nx6 (x1, y1, x2, y2, conf, cls) """ nc = prediction[0].shape[1] - 5 # number of classes xc = prediction[..., 4] > conf_thres # candidates # Settings min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height max_det = 300 # maximum number of detections per image time_limit = 10.0 # seconds to quit after redundant = True # require redundant detections multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img) t = time.time() output = [None] * prediction.shape[0] for xi, x in enumerate(prediction): # image index, image inference # Apply constraints # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height x = x[xc[xi]] # confidence # If none remain process next image if not x.shape[0]: continue # Compute conf x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf # Box (center x, center y, width, height) to (x1, y1, x2, y2) box = xywh2xyxy(x[:, :4]) # Detections matrix nx6 (xyxy, conf, cls) if multi_label: i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1) else: # best class only conf, j = x[:, 5:].max(1, keepdim=True) x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres] # Filter by class if classes: x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)] # Apply finite constraint # if not torch.isfinite(x).all(): # x = x[torch.isfinite(x).all(1)] # If none remain process next image n = x.shape[0] # number of boxes if not n: continue # Sort by confidence # x = x[x[:, 4].argsort(descending=True)] # Batched NMS c = x[:, 5:6] * (0 if agnostic else max_wh) # classes boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores i = torch.ops.torchvision.nms(boxes, scores, iou_thres) if i.shape[0] > max_det: # limit detections i = i[:max_det] if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean) try: # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4) iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix weights = iou * scores[None] # box weights x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes if redundant: i = i[iou.sum(1) > 1] # require redundancy except: # possible CUDA error https://github.com/ultralytics/yolov3/issues/1139 print(x, i, x.shape, i.shape) pass output[xi] = x[i] if (time.time() - t) > time_limit: break # time limit exceeded return output
原文：Whiten/customize the background of images using Semantic Segmentation and Alpha Blending - 2019.12.26作者：Shilpi Bhattacharyya图：语义分割和 alpha blending 技术用于改变图像的背景.图：使用语义分割和 alpha blending 的处理流程. 将背景区域像素值设为 0.完整实现：https://github.com/shilpibhattacharyya/customizing_image_backgrounds是否有使用 Photoshop 进行图像背景白化(whiten)和改变时抓狂的时候呢？有的话，请进入语义分割.1. 语义分割语义分割是一种将图像中每个像素进行归类的图像分析任务. 类似于人的方式，当人在看一幅图像时，知道图像的每个像素属于哪一类.更多关于语义分割的资料，可参考：Image Segmentation - 2018.11.05这里采用 DeepLabV3 进行语义分割，alpha blending 进行将原图融合到新背景.图：输入图像采用 Deeplab 模型语义分割处理. 期望输出是，覆盖在输入图像上的语义标签. 语义分割模型关注与对图像中每个像素分配语义标签，如sky, person, car 等.DeepLab 是 Google Brain 提出的语义分割结构，其是一种先进的语义分割深度学习模型. 其系列主要有：DeepLabv1 - 采用 atrous conv 来显式地控制网络中所计算的特征图的分辨率；DeepLabv2 - 采用 atrous spatial pyramid pooling(ASPP) 来鲁棒的分割物体，其在不同的采样率和有效接受野进行多尺度处理(using atrous spatial pyramid pooling (ASPP) to robustly segment objects at multiple scales with filters at multiple sampling rates and effective fields-of-views)；DeepLabv3 - 采用图像层(image-level)特征来增强 ASPP 模块，以捕获更大范围的特征. 此外，还包含 batch normalization 参数来提升模型训练. 具体来说，在训练和验证阶段，采用 atrous conv 提取不同输出步长(output strides)的输出特征，有效的增强了在步长=16 的 BN 训练，且验证阶段在步长=8时获得了更好的效果.google - deeplab2. 语义分割实现2.1. 分割输出解码函数：将 2D 分割图像转换为 RGB 图像，其中每一个标签被映射到对应的颜色.import numpy as np def decode_segmap(image, nc=21): label_colors = np.array([(0, 0, 0), # 0=background # 1=aeroplane, 2=bicycle, 3=bird, 4=boat, 5=bottle (128, 0, 0), (0, 128, 0), (255, 255, 255), (0, 0, 128), (128, 0, 128), # 6=bus, 7=car, 8=cat, 9=chair, 10=cow (0, 128, 128), (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0), # 11=dining table, 12=dog, 13=horse, 14=motorbike, 15=person (192, 128, 0), (64, 0, 128), (192, 0, 128), (64, 128, 128), (255, 255, 255), # 16=potted plant, 17=sheep, 18=sofa, 19=train, 20=tv/monitor (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128)]) r = np.zeros_like(image).astype(np.uint8) g = np.zeros_like(image).astype(np.uint8) b = np.zeros_like(image).astype(np.uint8) for l in range(0, nc): idx = image == l r[idx] = label_colors[l, 0] g[idx] = label_colors[l, 1] b[idx] = label_colors[l, 2] rgb = np.stack([r, g, b], axis=2) return rgb2.2. 图像预处理与Forwardimport matplotlib.pyplot as plt import torchvision.transforms as T import torch from PIL import Image def segment(net, path, show_orig=True, dev='cuda'): img = Image.open(path) if show_orig: plt.imshow(img); plt.axis('off'); plt.show() # Comment the Resize and CenterCrop for better inference results trf = T.Compose([T.ToTensor(), T.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225])]) inp = trf(img).unsqueeze(0).to(dev) out = net.to(dev)(inp)['out'] om = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy() rgb = decode_segmap(om) return rgb2.3. 加载 deeplabv3_resnet101 模型from torchvision import models dlab = models.segmentation.deeplabv3_resnet101(pretrained=1).eval() rgb=segment(dlab, '/path/to/test.jpg') ## If there are multiple labeled objects in the image, use the below code to have only the target as the foreground rgb[rgb!=255]=03. 背景融合(Alpha blending)实现Alpha blending 是将前景图像和透明图覆盖到背景图像上的处理. 透明图通常是图像的第四个通道(如，透明 PNG)，但其也可以是单独的图像. 透明 mask 图也一般被称为 alpha mask 或 alpha matte. Alpha blending 背后的数据原理是比较直接的，对于图像的每个像素，需要采用 alpha mask 将前景图像颜色和背景图像颜色组合在一起.前景图像是原始图像，灰度 alpha mask 是有语义分割得到的结果，背景图像是需要更换的新的定制图像.#alpha blending to customize the background of the image # Read the images foreground = cv2.imread("/path/to/foreground.jpg") background = cv2.imread("/path/to/background.jpg", cv2.IMREAD_COLOR) background = cv2.resize(background, (rgb.shape[1],rgb.shape[0]), interpolation = cv2.INTER_AREA) alpha = rgb # 2.3 # Convert uint8 to float foreground = foreground.astype(float) background = background.astype(float) # Normalize the alpha mask to keep intensity between 0 and 1 alpha = alpha.astype(float)/255 # Multiply the foreground with the alpha matte foreground = cv2.multiply(alpha, foreground) # Multiply the background with ( 1 - alpha ) background = cv2.multiply(1.0 - alpha, background) # Add the masked foreground and background. outImage = cv2.add(foreground, background) # Display image cv2_imshow(outImage) cv2.waitKey(0) # Save/download image cv2.imwrite('org_plus_cust_bkg_img.png' , outImage)如：图：语义分割和 alpha blending 的背景替换结果图：语义分割和 alpha blending 的处理过程4. 背景白化(Whiten)#whiten the background of the image mask_out=cv2.subtract(rgb,img) mask_out=cv2.subtract(rgb,mask_out) mask_out[rgb == 0] = 255 # Display the result numpy_horizontal = np.hstack((img, mask_out)) numpy_horizontal_concat = np.concatenate((img, mask_out), axis=1) cv2_imshow(numpy_horizontal_concat) cv2.waitKey(0) # Save/download the resulting image cv2.imwrite('org_plus_white_bkg_image.jpeg',numpy_horizontal_concat)如：图：图像背景白化结果5. 背景模糊(Blur)# Read the images foreground = cv2.imread("/path/to/girl8.jpg") # Create a Gaussian blur of kernel size 7 for the background image blurredImage = cv2.GaussianBlur(foreground, (7,7), 0) # Convert uint8 to float foreground = foreground.astype(float) blurredImage = blurredImage.astype(float) # Create a binary mask of the RGB output map using the threshold value 0 th, alpha = cv2.threshold(np.array(rgb),0, 255, cv2.THRESH_BINARY) # Apply a slight blur to the mask to soften edges alpha = cv2.GaussianBlur(alpha, (7,7),0) # Normalize the alpha mask to keep intensity between 0 and 1 alpha = alpha.astype(float)/255 # Multiply the foreground with the alpha matte foreground = cv2.multiply(alpha, foreground) # Multiply the background with ( 1 - alpha ) background = cv2.multiply(1.0 - alpha, blurredImage) # Add the masked foreground and background outImage = cv2.add(foreground, background) # Return a normalized output image for display outImage= outImage numpy_horizontal = np.hstack((img, outImage)) numpy_horizontal_concat = np.concatenate((img, outImage), axis=1) # Display image cv2_imshow(numpy_horizontal_concat) cv2.waitKey(0) # Save/download the resulting image cv2.imwrite('res.png' , numpy_horizontal_concat)如：图：图像背景融合结果6. 背景灰度化(Grayscale)# Load the foreground input image foreground = cv2.imread("/path/to/woman-with-car.jpg") # Resize image to match shape of R-band in RGB output map foreground = cv2.resize(foreground, (rgb.shape[1],rgb.shape[0]), interpolation = cv2.INTER_AREA) # Create a background image by copying foreground and converting into grayscale background = cv2.cvtColor(foreground, cv2.COLOR_BGR2GRAY) # convert single channel grayscale image to 3-channel grayscale image background = cv2.cvtColor(background, cv2.COLOR_GRAY2RGB) # Convert uint8 to float foreground = foreground.astype(float) background = background.astype(float) # Create a binary mask of the RGB output map using the threshold value 0 th, alpha = cv2.threshold(np.array(rgb),0,255, cv2.THRESH_BINARY) # Apply a slight blur to the mask to soften edges alpha = cv2.GaussianBlur(alpha, (7,7),0) # Normalize the alpha mask to keep intensity between 0 and 1 alpha = alpha.astype(float)/255 # Multiply the foreground with the alpha matte foreground = cv2.multiply(alpha, foreground) # Multiply the background with ( 1 - alpha ) background = cv2.multiply(1.0 - alpha, background) # Add the masked foreground and background outImage = cv2.add(foreground, background) numpy_horizontal = np.hstack((img, outImage)) numpy_horizontal_concat = np.concatenate((img, outImage), axis=1) # Display image cv2_imshow(numpy_horizontal_concat) cv2.waitKey(0) # Save image cv2.imwrite('res.png' , numpy_horizontal_concat)如：图：背景灰度化结果图：背景灰度化处理过程参考Alpha Blending — https://www.learnopencv.com/alpha-blending-using-opencv-cpp-python/Semantic Segmentation — https://www.learnopencv.com/pytorch-for-beginners-semantic-segmentation-using-torchvision/DeepLab — https://github.com/tensorflow/models/tree/master/research/deeplabhttps://www.learnopencv.com/applications-of-foreground-background-separation-with-semantic-segmentation/