From 802f90d2e1a73188e821ea832b05bab7967e4f1e Mon Sep 17 00:00:00 2001 From: pipikk Date: Mon, 8 Apr 2024 15:27:20 +0800 Subject: [PATCH] feat: add raft model (#702) * first version * fix minor bugs * add test images * TorchModel & docstr * modify pipeline implementation * minor * add datasets * add extractor * add update * add augmentor * add flow_viz * add frame_utils * add utils * add raft_model * add dense_optical_flow_estimation_pipeline * add image_utils * add test_dense_optical_flow_estimation * test * update cv/__init__ * [3]update cv/__init__ * update submodule data/test * correct yapf * fix bugs * move test data * update submodule * minor * update submodule --------- Co-authored-by: kejie --- data/test | 2 +- modelscope/metainfo.py | 11 +- modelscope/models/cv/__init__.py | 10 +- .../dense_optical_flow_estimation/__init__.py | 21 ++ .../core/__init__.py | 0 .../core/corr.py | 95 ++++++ .../core/datasets.py | 297 ++++++++++++++++++ .../core/extractor.py | 285 +++++++++++++++++ .../core/raft.py | 163 ++++++++++ .../core/update.py | 157 +++++++++ .../core/utils/__init__.py | 0 .../core/utils/augmentor.py | 286 +++++++++++++++++ .../core/utils/flow_viz.py | 132 ++++++++ .../core/utils/frame_utils.py | 142 +++++++++ .../core/utils/utils.py | 93 ++++++ .../raft_model.py | 52 +++ modelscope/outputs/outputs.py | 2 + .../dense_optical_flow_estimation_pipeline.py | 147 +++++++++ modelscope/utils/constant.py | 1 + modelscope/utils/cv/image_utils.py | 146 +++++++++ modelscope/utils/pipeline_schema.json | 7 + .../test_dense_optical_flow_estimation.py | 39 +++ 22 files changed, 2079 insertions(+), 9 deletions(-) create mode 100644 modelscope/models/cv/dense_optical_flow_estimation/__init__.py create mode 100644 modelscope/models/cv/dense_optical_flow_estimation/core/__init__.py create mode 100644 modelscope/models/cv/dense_optical_flow_estimation/core/corr.py create mode 100644 modelscope/models/cv/dense_optical_flow_estimation/core/datasets.py create mode 100644 modelscope/models/cv/dense_optical_flow_estimation/core/extractor.py create mode 100644 modelscope/models/cv/dense_optical_flow_estimation/core/raft.py create mode 100644 modelscope/models/cv/dense_optical_flow_estimation/core/update.py create mode 100644 modelscope/models/cv/dense_optical_flow_estimation/core/utils/__init__.py create mode 100644 modelscope/models/cv/dense_optical_flow_estimation/core/utils/augmentor.py create mode 100644 modelscope/models/cv/dense_optical_flow_estimation/core/utils/flow_viz.py create mode 100644 modelscope/models/cv/dense_optical_flow_estimation/core/utils/frame_utils.py create mode 100644 modelscope/models/cv/dense_optical_flow_estimation/core/utils/utils.py create mode 100644 modelscope/models/cv/dense_optical_flow_estimation/raft_model.py create mode 100644 modelscope/pipelines/cv/dense_optical_flow_estimation_pipeline.py create mode 100644 tests/pipelines/test_dense_optical_flow_estimation.py diff --git a/data/test b/data/test index 860764da..7a7f6b8d 160000 --- a/data/test +++ b/data/test @@ -1 +1 @@ -Subproject commit 860764da23420f08fa551eccc053719b8f1a4b42 +Subproject commit 7a7f6b8d05ba8af4ea42096391fa727d358e585e diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 772dbb28..00d61c8b 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -57,6 +57,7 @@ class Models(object): unifuse_depth_estimation = 'unifuse-depth-estimation' s2net_depth_estimation = 's2net-depth-estimation' dro_resnet18_depth_estimation = 'dro-resnet18-depth-estimation' + raft_dense_optical_flow_estimation = 'raft-dense-optical-flow-estimation' resnet50_bert = 'resnet50-bert' referring_video_object_segmentation = 'swinT-referring-video-object-segmentation' fer = 'fer' @@ -404,6 +405,7 @@ class Pipelines(object): video_depth_estimation = 'video-depth-estimation' panorama_depth_estimation = 'panorama-depth-estimation' panorama_depth_estimation_s2net = 'panorama-depth-estimation-s2net' + dense_optical_flow_estimation = 'dense-optical-flow-estimation' image_reid_person = 'passvitb-image-reid-person' image_inpainting = 'fft-inpainting' image_paintbyexample = 'stablediffusion-paintbyexample' @@ -815,6 +817,9 @@ DEFAULT_MODEL_FOR_PIPELINE = { Tasks.panorama_depth_estimation: (Pipelines.panorama_depth_estimation, 'damo/cv_unifuse_panorama-depth-estimation'), + Tasks.dense_optical_flow_estimation: + (Pipelines.dense_optical_flow_estimation, + 'Damo_XR_Lab/cv_raft_dense-optical-flow_things'), Tasks.image_local_feature_matching: (Pipelines.image_local_feature_matching, 'Damo_XR_Lab/cv_resnet-transformer_local-feature-matching_outdoor-data'), @@ -838,9 +843,9 @@ DEFAULT_MODEL_FOR_PIPELINE = { Tasks.image_classification: (Pipelines.daily_image_classification, 'damo/cv_vit-base_image-classification_Dailylife-labels'), - Tasks.image_object_detection: - (Pipelines.image_object_detection_auto, - 'damo/cv_yolox_image-object-detection-auto'), + Tasks.image_object_detection: ( + Pipelines.image_object_detection_auto, + 'damo/cv_yolox_image-object-detection-auto'), Tasks.ocr_recognition: ( Pipelines.ocr_recognition, 'damo/cv_convnextTiny_ocr-recognition-general_damo'), diff --git a/modelscope/models/cv/__init__.py b/modelscope/models/cv/__init__.py index 52da23b8..2bf632f8 100644 --- a/modelscope/models/cv/__init__.py +++ b/modelscope/models/cv/__init__.py @@ -4,11 +4,11 @@ from . import (action_recognition, animal_recognition, bad_image_detecting, body_2d_keypoints, body_3d_keypoints, cartoon, cmdssl_video_embedding, controllable_image_generation, - crowd_counting, face_detection, face_generation, - face_reconstruction, human3d_animation, human_reconstruction, - image_classification, image_color_enhance, image_colorization, - image_defrcn_fewshot, image_denoise, image_editing, - image_inpainting, image_instance_segmentation, + crowd_counting, dense_optical_flow_estimation, face_detection, + face_generation, face_reconstruction, human3d_animation, + human_reconstruction, image_classification, image_color_enhance, + image_colorization, image_defrcn_fewshot, image_denoise, + image_editing, image_inpainting, image_instance_segmentation, image_local_feature_matching, image_matching, image_matching_fast, image_mvs_depth_estimation, image_mvs_depth_estimation_geomvsnet, diff --git a/modelscope/models/cv/dense_optical_flow_estimation/__init__.py b/modelscope/models/cv/dense_optical_flow_estimation/__init__.py new file mode 100644 index 00000000..be8fc28e --- /dev/null +++ b/modelscope/models/cv/dense_optical_flow_estimation/__init__.py @@ -0,0 +1,21 @@ +from typing import TYPE_CHECKING + +from modelscope.utils.import_utils import LazyImportModule + +if TYPE_CHECKING: + from .raft_model import DenseOpticalFlowEstimation + +else: + _import_structure = { + 'raft_dense_optical_flow_estimation': ['DenseOpticalFlowEstimation'], + } + + import sys + + sys.modules[__name__] = LazyImportModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/modelscope/models/cv/dense_optical_flow_estimation/core/__init__.py b/modelscope/models/cv/dense_optical_flow_estimation/core/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/dense_optical_flow_estimation/core/corr.py b/modelscope/models/cv/dense_optical_flow_estimation/core/corr.py new file mode 100644 index 00000000..a0b1a27e --- /dev/null +++ b/modelscope/models/cv/dense_optical_flow_estimation/core/corr.py @@ -0,0 +1,95 @@ +import torch +import torch.nn.functional as F + +from modelscope.models.cv.dense_optical_flow_estimation.core.utils.utils import ( + bilinear_sampler, coords_grid) + +try: + import alt_cuda_corr +except ModuleNotFoundError: + # alt_cuda_corr is not compiled + pass + + +class CorrBlock: + + def __init__(self, fmap1, fmap2, num_levels=4, radius=4): + self.num_levels = num_levels + self.radius = radius + self.corr_pyramid = [] + + # all pairs correlation + corr = CorrBlock.corr(fmap1, fmap2) + + batch, h1, w1, dim, h2, w2 = corr.shape + corr = corr.reshape(batch * h1 * w1, dim, h2, w2) + + self.corr_pyramid.append(corr) + for i in range(self.num_levels - 1): + corr = F.avg_pool2d(corr, 2, stride=2) + self.corr_pyramid.append(corr) + + def __call__(self, coords): + r = self.radius + coords = coords.permute(0, 2, 3, 1) + batch, h1, w1, _ = coords.shape + + out_pyramid = [] + for i in range(self.num_levels): + corr = self.corr_pyramid[i] + dx = torch.linspace(-r, r, 2 * r + 1, device=coords.device) + dy = torch.linspace(-r, r, 2 * r + 1, device=coords.device) + delta = torch.stack(torch.meshgrid(dy, dx), axis=-1) + + centroid_lvl = coords.reshape(batch * h1 * w1, 1, 1, 2) / 2**i + delta_lvl = delta.view(1, 2 * r + 1, 2 * r + 1, 2) + coords_lvl = centroid_lvl + delta_lvl + + corr = bilinear_sampler(corr, coords_lvl) + corr = corr.view(batch, h1, w1, -1) + out_pyramid.append(corr) + + out = torch.cat(out_pyramid, dim=-1) + return out.permute(0, 3, 1, 2).contiguous().float() + + @staticmethod + def corr(fmap1, fmap2): + batch, dim, ht, wd = fmap1.shape + fmap1 = fmap1.view(batch, dim, ht * wd) + fmap2 = fmap2.view(batch, dim, ht * wd) + + corr = torch.matmul(fmap1.transpose(1, 2), fmap2) + corr = corr.view(batch, ht, wd, 1, ht, wd) + return corr / torch.sqrt(torch.tensor(dim).float()) + + +class AlternateCorrBlock: + + def __init__(self, fmap1, fmap2, num_levels=4, radius=4): + self.num_levels = num_levels + self.radius = radius + + self.pyramid = [(fmap1, fmap2)] + for i in range(self.num_levels): + fmap1 = F.avg_pool2d(fmap1, 2, stride=2) + fmap2 = F.avg_pool2d(fmap2, 2, stride=2) + self.pyramid.append((fmap1, fmap2)) + + def __call__(self, coords): + coords = coords.permute(0, 2, 3, 1) + B, H, W, _ = coords.shape + dim = self.pyramid[0][0].shape[1] + + corr_list = [] + for i in range(self.num_levels): + r = self.radius + fmap1_i = self.pyramid[0][0].permute(0, 2, 3, 1).contiguous() + fmap2_i = self.pyramid[i][1].permute(0, 2, 3, 1).contiguous() + + coords_i = (coords / 2**i).reshape(B, 1, H, W, 2).contiguous() + corr, = alt_cuda_corr.forward(fmap1_i, fmap2_i, coords_i, r) + corr_list.append(corr.squeeze(1)) + + corr = torch.stack(corr_list, dim=1) + corr = corr.reshape(B, -1, H, W) + return corr / torch.sqrt(torch.tensor(dim).float()) diff --git a/modelscope/models/cv/dense_optical_flow_estimation/core/datasets.py b/modelscope/models/cv/dense_optical_flow_estimation/core/datasets.py new file mode 100644 index 00000000..eb8a8559 --- /dev/null +++ b/modelscope/models/cv/dense_optical_flow_estimation/core/datasets.py @@ -0,0 +1,297 @@ +# Data loading based on https://github.com/NVIDIA/flownet2-pytorch + +import math +import os +import os.path as osp +import random +from glob import glob + +import numpy as np +import torch +import torch.nn.functional as F +import torch.utils.data as data +from utils import frame_utils +from utils.augmentor import FlowAugmentor, SparseFlowAugmentor + + +class FlowDataset(data.Dataset): + + def __init__(self, aug_params=None, sparse=False): + self.augmentor = None + self.sparse = sparse + if aug_params is not None: + if sparse: + self.augmentor = SparseFlowAugmentor(**aug_params) + else: + self.augmentor = FlowAugmentor(**aug_params) + + self.is_test = False + self.init_seed = False + self.flow_list = [] + self.image_list = [] + self.extra_info = [] + + def __getitem__(self, index): + + if self.is_test: + img1 = frame_utils.read_gen(self.image_list[index][0]) + img2 = frame_utils.read_gen(self.image_list[index][1]) + img1 = np.array(img1).astype(np.uint8)[..., :3] + img2 = np.array(img2).astype(np.uint8)[..., :3] + img1 = torch.from_numpy(img1).permute(2, 0, 1).float() + img2 = torch.from_numpy(img2).permute(2, 0, 1).float() + return img1, img2, self.extra_info[index] + + if not self.init_seed: + worker_info = torch.utils.data.get_worker_info() + if worker_info is not None: + torch.manual_seed(worker_info.id) + np.random.seed(worker_info.id) + random.seed(worker_info.id) + self.init_seed = True + + index = index % len(self.image_list) + valid = None + if self.sparse: + flow, valid = frame_utils.readFlowKITTI(self.flow_list[index]) + else: + flow = frame_utils.read_gen(self.flow_list[index]) + + img1 = frame_utils.read_gen(self.image_list[index][0]) + img2 = frame_utils.read_gen(self.image_list[index][1]) + + flow = np.array(flow).astype(np.float32) + img1 = np.array(img1).astype(np.uint8) + img2 = np.array(img2).astype(np.uint8) + + # grayscale images + if len(img1.shape) == 2: + img1 = np.tile(img1[..., None], (1, 1, 3)) + img2 = np.tile(img2[..., None], (1, 1, 3)) + else: + img1 = img1[..., :3] + img2 = img2[..., :3] + + if self.augmentor is not None: + if self.sparse: + img1, img2, flow, valid = self.augmentor( + img1, img2, flow, valid) + else: + img1, img2, flow = self.augmentor(img1, img2, flow) + + img1 = torch.from_numpy(img1).permute(2, 0, 1).float() + img2 = torch.from_numpy(img2).permute(2, 0, 1).float() + flow = torch.from_numpy(flow).permute(2, 0, 1).float() + + if valid is not None: + valid = torch.from_numpy(valid) + else: + valid = (flow[0].abs() < 1000) & (flow[1].abs() < 1000) + + return img1, img2, flow, valid.float() + + def __rmul__(self, v): + self.flow_list = v * self.flow_list + self.image_list = v * self.image_list + return self + + def __len__(self): + return len(self.image_list) + + +class MpiSintel(FlowDataset): + + def __init__(self, + aug_params=None, + split='training', + root='datasets/Sintel', + dstype='clean'): + super(MpiSintel, self).__init__(aug_params) + flow_root = osp.join(root, split, 'flow') + image_root = osp.join(root, split, dstype) + + if split == 'test': + self.is_test = True + + for scene in os.listdir(image_root): + image_list = sorted(glob(osp.join(image_root, scene, '*.png'))) + for i in range(len(image_list) - 1): + self.image_list += [[image_list[i], image_list[i + 1]]] + self.extra_info += [(scene, i)] # scene and frame_id + + if split != 'test': + self.flow_list += sorted( + glob(osp.join(flow_root, scene, '*.flo'))) + + +class FlyingChairs(FlowDataset): + + def __init__(self, + aug_params=None, + split='train', + root='datasets/FlyingChairs_release/data'): + super(FlyingChairs, self).__init__(aug_params) + + images = sorted(glob(osp.join(root, '*.ppm'))) + flows = sorted(glob(osp.join(root, '*.flo'))) + assert (len(images) // 2 == len(flows)) + + split_list = np.loadtxt('chairs_split.txt', dtype=np.int32) + for i in range(len(flows)): + xid = split_list[i] + if (split == 'training' and xid == 1) or (split == 'validation' + and xid == 2): + self.flow_list += [flows[i]] + self.image_list += [[images[2 * i], images[2 * i + 1]]] + + +class FlyingThings3D(FlowDataset): + + def __init__(self, + aug_params=None, + root='datasets/FlyingThings3D', + dstype='frames_cleanpass'): + super(FlyingThings3D, self).__init__(aug_params) + + for cam in ['left']: + for direction in ['into_future', 'into_past']: + image_dirs = sorted(glob(osp.join(root, dstype, 'TRAIN/*/*'))) + image_dirs = sorted([osp.join(f, cam) for f in image_dirs]) + + flow_dirs = sorted( + glob(osp.join(root, 'optical_flow/TRAIN/*/*'))) + flow_dirs = sorted( + [osp.join(f, direction, cam) for f in flow_dirs]) + + for idir, fdir in zip(image_dirs, flow_dirs): + images = sorted(glob(osp.join(idir, '*.png'))) + flows = sorted(glob(osp.join(fdir, '*.pfm'))) + for i in range(len(flows) - 1): + if direction == 'into_future': + self.image_list += [[images[i], images[i + 1]]] + self.flow_list += [flows[i]] + elif direction == 'into_past': + self.image_list += [[images[i + 1], images[i]]] + self.flow_list += [flows[i + 1]] + + +class KITTI(FlowDataset): + + def __init__(self, + aug_params=None, + split='training', + root='datasets/KITTI'): + super(KITTI, self).__init__(aug_params, sparse=True) + if split == 'testing': + self.is_test = True + + root = osp.join(root, split) + images1 = sorted(glob(osp.join(root, 'image_2/*_10.png'))) + images2 = sorted(glob(osp.join(root, 'image_2/*_11.png'))) + + for img1, img2 in zip(images1, images2): + frame_id = img1.split('/')[-1] + self.extra_info += [[frame_id]] + self.image_list += [[img1, img2]] + + if split == 'training': + self.flow_list = sorted(glob(osp.join(root, 'flow_occ/*_10.png'))) + + +class HD1K(FlowDataset): + + def __init__(self, aug_params=None, root='datasets/HD1k'): + super(HD1K, self).__init__(aug_params, sparse=True) + + seq_ix = 0 + while 1: + flows = sorted( + glob( + os.path.join(root, 'hd1k_flow_gt', + 'flow_occ/%06d_*.png' % seq_ix))) + images = sorted( + glob( + os.path.join(root, 'hd1k_input', + 'image_2/%06d_*.png' % seq_ix))) + + if len(flows) == 0: + break + + for i in range(len(flows) - 1): + self.flow_list += [flows[i]] + self.image_list += [[images[i], images[i + 1]]] + + seq_ix += 1 + + +def fetch_dataloader(args, TRAIN_DS='C+T+K+S+H'): + """ Create the data loader for the corresponding trainign set """ + + if args.stage == 'chairs': + aug_params = { + 'crop_size': args.image_size, + 'min_scale': -0.1, + 'max_scale': 1.0, + 'do_flip': True + } + train_dataset = FlyingChairs(aug_params, split='training') + + elif args.stage == 'things': + aug_params = { + 'crop_size': args.image_size, + 'min_scale': -0.4, + 'max_scale': 0.8, + 'do_flip': True + } + clean_dataset = FlyingThings3D(aug_params, dstype='frames_cleanpass') + final_dataset = FlyingThings3D(aug_params, dstype='frames_finalpass') + train_dataset = clean_dataset + final_dataset + + elif args.stage == 'sintel': + aug_params = { + 'crop_size': args.image_size, + 'min_scale': -0.2, + 'max_scale': 0.6, + 'do_flip': True + } + things = FlyingThings3D(aug_params, dstype='frames_cleanpass') + sintel_clean = MpiSintel(aug_params, split='training', dstype='clean') + sintel_final = MpiSintel(aug_params, split='training', dstype='final') + + if TRAIN_DS == 'C+T+K+S+H': + kitti = KITTI({ + 'crop_size': args.image_size, + 'min_scale': -0.3, + 'max_scale': 0.5, + 'do_flip': True + }) + hd1k = HD1K({ + 'crop_size': args.image_size, + 'min_scale': -0.5, + 'max_scale': 0.2, + 'do_flip': True + }) + train_dataset = 100 * sintel_clean + 100 * sintel_final + 200 * kitti + 5 * hd1k + things + + elif TRAIN_DS == 'C+T+K/S': + train_dataset = 100 * sintel_clean + 100 * sintel_final + things + + elif args.stage == 'kitti': + aug_params = { + 'crop_size': args.image_size, + 'min_scale': -0.2, + 'max_scale': 0.4, + 'do_flip': False + } + train_dataset = KITTI(aug_params, split='training') + + train_loader = data.DataLoader( + train_dataset, + batch_size=args.batch_size, + pin_memory=False, + shuffle=True, + num_workers=4, + drop_last=True) + + print('Training with %d image pairs' % len(train_dataset)) + return train_loader diff --git a/modelscope/models/cv/dense_optical_flow_estimation/core/extractor.py b/modelscope/models/cv/dense_optical_flow_estimation/core/extractor.py new file mode 100644 index 00000000..dfa8e4de --- /dev/null +++ b/modelscope/models/cv/dense_optical_flow_estimation/core/extractor.py @@ -0,0 +1,285 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class ResidualBlock(nn.Module): + + def __init__(self, in_planes, planes, norm_fn='group', stride=1): + super(ResidualBlock, self).__init__() + + self.conv1 = nn.Conv2d( + in_planes, planes, kernel_size=3, padding=1, stride=stride) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1) + self.relu = nn.ReLU(inplace=True) + + num_groups = planes // 8 + + if norm_fn == 'group': + self.norm1 = nn.GroupNorm( + num_groups=num_groups, num_channels=planes) + self.norm2 = nn.GroupNorm( + num_groups=num_groups, num_channels=planes) + if not stride == 1: + self.norm3 = nn.GroupNorm( + num_groups=num_groups, num_channels=planes) + + elif norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(planes) + self.norm2 = nn.BatchNorm2d(planes) + if not stride == 1: + self.norm3 = nn.BatchNorm2d(planes) + + elif norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(planes) + self.norm2 = nn.InstanceNorm2d(planes) + if not stride == 1: + self.norm3 = nn.InstanceNorm2d(planes) + + elif norm_fn == 'none': + self.norm1 = nn.Sequential() + self.norm2 = nn.Sequential() + if not stride == 1: + self.norm3 = nn.Sequential() + + if stride == 1: + self.downsample = None + + else: + self.downsample = nn.Sequential( + nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), + self.norm3) + + def forward(self, x): + y = x + y = self.relu(self.norm1(self.conv1(y))) + y = self.relu(self.norm2(self.conv2(y))) + + if self.downsample is not None: + x = self.downsample(x) + + return self.relu(x + y) + + +class BottleneckBlock(nn.Module): + + def __init__(self, in_planes, planes, norm_fn='group', stride=1): + super(BottleneckBlock, self).__init__() + + self.conv1 = nn.Conv2d( + in_planes, planes // 4, kernel_size=1, padding=0) + self.conv2 = nn.Conv2d( + planes // 4, planes // 4, kernel_size=3, padding=1, stride=stride) + self.conv3 = nn.Conv2d(planes // 4, planes, kernel_size=1, padding=0) + self.relu = nn.ReLU(inplace=True) + + num_groups = planes // 8 + + if norm_fn == 'group': + self.norm1 = nn.GroupNorm( + num_groups=num_groups, num_channels=planes // 4) + self.norm2 = nn.GroupNorm( + num_groups=num_groups, num_channels=planes // 4) + self.norm3 = nn.GroupNorm( + num_groups=num_groups, num_channels=planes) + if not stride == 1: + self.norm4 = nn.GroupNorm( + num_groups=num_groups, num_channels=planes) + + elif norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(planes // 4) + self.norm2 = nn.BatchNorm2d(planes // 4) + self.norm3 = nn.BatchNorm2d(planes) + if not stride == 1: + self.norm4 = nn.BatchNorm2d(planes) + + elif norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(planes // 4) + self.norm2 = nn.InstanceNorm2d(planes // 4) + self.norm3 = nn.InstanceNorm2d(planes) + if not stride == 1: + self.norm4 = nn.InstanceNorm2d(planes) + + elif norm_fn == 'none': + self.norm1 = nn.Sequential() + self.norm2 = nn.Sequential() + self.norm3 = nn.Sequential() + if not stride == 1: + self.norm4 = nn.Sequential() + + if stride == 1: + self.downsample = None + + else: + self.downsample = nn.Sequential( + nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), + self.norm4) + + def forward(self, x): + y = x + y = self.relu(self.norm1(self.conv1(y))) + y = self.relu(self.norm2(self.conv2(y))) + y = self.relu(self.norm3(self.conv3(y))) + + if self.downsample is not None: + x = self.downsample(x) + + return self.relu(x + y) + + +class BasicEncoder(nn.Module): + + def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0): + super(BasicEncoder, self).__init__() + self.norm_fn = norm_fn + + if self.norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64) + + elif self.norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(64) + + elif self.norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(64) + + elif self.norm_fn == 'none': + self.norm1 = nn.Sequential() + + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3) + self.relu1 = nn.ReLU(inplace=True) + + self.in_planes = 64 + self.layer1 = self._make_layer(64, stride=1) + self.layer2 = self._make_layer(96, stride=2) + self.layer3 = self._make_layer(128, stride=2) + + # output convolution + self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1) + + self.dropout = None + if dropout > 0: + self.dropout = nn.Dropout2d(p=dropout) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_( + m.weight, mode='fan_out', nonlinearity='relu') + elif isinstance(m, + (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): + if m.weight is not None: + nn.init.constant_(m.weight, 1) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def _make_layer(self, dim, stride=1): + layer1 = ResidualBlock( + self.in_planes, dim, self.norm_fn, stride=stride) + layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1) + layers = (layer1, layer2) + + self.in_planes = dim + return nn.Sequential(*layers) + + def forward(self, x): + + # if input is list, combine batch dimension + is_list = isinstance(x, tuple) or isinstance(x, list) + if is_list: + batch_dim = x[0].shape[0] + x = torch.cat(x, dim=0) + + x = self.conv1(x) + x = self.norm1(x) + x = self.relu1(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + + x = self.conv2(x) + + if self.training and self.dropout is not None: + x = self.dropout(x) + + if is_list: + x = torch.split(x, [batch_dim, batch_dim], dim=0) + + return x + + +class SmallEncoder(nn.Module): + + def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0): + super(SmallEncoder, self).__init__() + self.norm_fn = norm_fn + + if self.norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=8, num_channels=32) + + elif self.norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(32) + + elif self.norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(32) + + elif self.norm_fn == 'none': + self.norm1 = nn.Sequential() + + self.conv1 = nn.Conv2d(3, 32, kernel_size=7, stride=2, padding=3) + self.relu1 = nn.ReLU(inplace=True) + + self.in_planes = 32 + self.layer1 = self._make_layer(32, stride=1) + self.layer2 = self._make_layer(64, stride=2) + self.layer3 = self._make_layer(96, stride=2) + + self.dropout = None + if dropout > 0: + self.dropout = nn.Dropout2d(p=dropout) + + self.conv2 = nn.Conv2d(96, output_dim, kernel_size=1) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_( + m.weight, mode='fan_out', nonlinearity='relu') + elif isinstance(m, + (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): + if m.weight is not None: + nn.init.constant_(m.weight, 1) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def _make_layer(self, dim, stride=1): + layer1 = BottleneckBlock( + self.in_planes, dim, self.norm_fn, stride=stride) + layer2 = BottleneckBlock(dim, dim, self.norm_fn, stride=1) + layers = (layer1, layer2) + + self.in_planes = dim + return nn.Sequential(*layers) + + def forward(self, x): + + # if input is list, combine batch dimension + is_list = isinstance(x, tuple) or isinstance(x, list) + if is_list: + batch_dim = x[0].shape[0] + x = torch.cat(x, dim=0) + + x = self.conv1(x) + x = self.norm1(x) + x = self.relu1(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.conv2(x) + + if self.training and self.dropout is not None: + x = self.dropout(x) + + if is_list: + x = torch.split(x, [batch_dim, batch_dim], dim=0) + + return x diff --git a/modelscope/models/cv/dense_optical_flow_estimation/core/raft.py b/modelscope/models/cv/dense_optical_flow_estimation/core/raft.py new file mode 100644 index 00000000..f2b801bc --- /dev/null +++ b/modelscope/models/cv/dense_optical_flow_estimation/core/raft.py @@ -0,0 +1,163 @@ +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from modelscope.models.base.base_torch_model import TorchModel +from modelscope.models.cv.dense_optical_flow_estimation.core.corr import ( + AlternateCorrBlock, CorrBlock) +from modelscope.models.cv.dense_optical_flow_estimation.core.extractor import ( + BasicEncoder, SmallEncoder) +from modelscope.models.cv.dense_optical_flow_estimation.core.update import ( + BasicUpdateBlock, SmallUpdateBlock) +from modelscope.models.cv.dense_optical_flow_estimation.core.utils.utils import ( + bilinear_sampler, coords_grid, upflow8) + +autocast = torch.cuda.amp.autocast + +# try: +# autocast = torch.cuda.amp.autocast +# except: +# # dummy autocast for PyTorch < 1.6 +# class autocast: +# def __init__(self, enabled): +# pass +# def __enter__(self): +# pass +# def __exit__(self, *args): +# pass + + +class RAFT(TorchModel): + + def __init__(self, args): + super(RAFT, self).__init__() + self.args = args + + if args.small: + self.hidden_dim = hdim = 96 + self.context_dim = cdim = 64 + args.corr_levels = 4 + args.corr_radius = 3 + + else: + self.hidden_dim = hdim = 128 + self.context_dim = cdim = 128 + args.corr_levels = 4 + args.corr_radius = 4 + + if 'dropout' not in self.args: + self.args.dropout = 0 + + if 'alternate_corr' not in self.args: + self.args.alternate_corr = False + + # feature network, context network, and update block + if args.small: + self.fnet = SmallEncoder( + output_dim=128, norm_fn='instance', dropout=args.dropout) + self.cnet = SmallEncoder( + output_dim=hdim + cdim, norm_fn='none', dropout=args.dropout) + self.update_block = SmallUpdateBlock(self.args, hidden_dim=hdim) + + else: + self.fnet = BasicEncoder( + output_dim=256, norm_fn='instance', dropout=args.dropout) + self.cnet = BasicEncoder( + output_dim=hdim + cdim, norm_fn='batch', dropout=args.dropout) + self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim) + + def freeze_bn(self): + for m in self.modules(): + if isinstance(m, nn.BatchNorm2d): + m.eval() + + def initialize_flow(self, img): + """ Flow is represented as difference between two coordinate grids flow = coords1 - coords0""" + N, C, H, W = img.shape + coords0 = coords_grid(N, H // 8, W // 8, device=img.device) + coords1 = coords_grid(N, H // 8, W // 8, device=img.device) + + # optical flow computed as difference: flow = coords1 - coords0 + return coords0, coords1 + + def upsample_flow(self, flow, mask): + """ Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """ + N, _, H, W = flow.shape + mask = mask.view(N, 1, 9, 8, 8, H, W) + mask = torch.softmax(mask, dim=2) + + up_flow = F.unfold(8 * flow, [3, 3], padding=1) + up_flow = up_flow.view(N, 2, 9, 1, 1, H, W) + + up_flow = torch.sum(mask * up_flow, dim=2) + up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) + return up_flow.reshape(N, 2, 8 * H, 8 * W) + + def forward(self, + image1, + image2, + iters=20, + flow_init=None, + upsample=True, + test_mode=False): + """ Estimate optical flow between pair of frames """ + + image1 = 2 * (image1 / 255.0) - 1.0 + image2 = 2 * (image2 / 255.0) - 1.0 + + image1 = image1.contiguous() + image2 = image2.contiguous() + + hdim = self.hidden_dim + cdim = self.context_dim + + # run the feature network + with autocast(enabled=self.args.mixed_precision): + fmap1, fmap2 = self.fnet([image1, image2]) + + fmap1 = fmap1.float() + fmap2 = fmap2.float() + if self.args.alternate_corr: + corr_fn = AlternateCorrBlock( + fmap1, fmap2, radius=self.args.corr_radius) + else: + corr_fn = CorrBlock(fmap1, fmap2, radius=self.args.corr_radius) + + # run the context network + with autocast(enabled=self.args.mixed_precision): + cnet = self.cnet(image1) + net, inp = torch.split(cnet, [hdim, cdim], dim=1) + net = torch.tanh(net) + inp = torch.relu(inp) + + coords0, coords1 = self.initialize_flow(image1) + + if flow_init is not None: + coords1 = coords1 + flow_init + + flow_predictions = [] + for itr in range(iters): + coords1 = coords1.detach() + corr = corr_fn(coords1) # index correlation volume + + flow = coords1 - coords0 + with autocast(enabled=self.args.mixed_precision): + net, up_mask, delta_flow = self.update_block( + net, inp, corr, flow) + + # F(t+1) = F(t) + \Delta(t) + coords1 = coords1 + delta_flow + + # upsample predictions + if up_mask is None: + flow_up = upflow8(coords1 - coords0) + else: + flow_up = self.upsample_flow(coords1 - coords0, up_mask) + + flow_predictions.append(flow_up) + + if test_mode: + return coords1 - coords0, flow_up + + return flow_predictions diff --git a/modelscope/models/cv/dense_optical_flow_estimation/core/update.py b/modelscope/models/cv/dense_optical_flow_estimation/core/update.py new file mode 100644 index 00000000..b43bb0ec --- /dev/null +++ b/modelscope/models/cv/dense_optical_flow_estimation/core/update.py @@ -0,0 +1,157 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class FlowHead(nn.Module): + + def __init__(self, input_dim=128, hidden_dim=256): + super(FlowHead, self).__init__() + self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1) + self.conv2 = nn.Conv2d(hidden_dim, 2, 3, padding=1) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + return self.conv2(self.relu(self.conv1(x))) + + +class ConvGRU(nn.Module): + + def __init__(self, hidden_dim=128, input_dim=192 + 128): + super(ConvGRU, self).__init__() + self.convz = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, 3, padding=1) + self.convr = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, 3, padding=1) + self.convq = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, 3, padding=1) + + def forward(self, h, x): + hx = torch.cat([h, x], dim=1) + + z = torch.sigmoid(self.convz(hx)) + r = torch.sigmoid(self.convr(hx)) + q = torch.tanh(self.convq(torch.cat([r * h, x], dim=1))) + + h = (1 - z) * h + z * q + return h + + +class SepConvGRU(nn.Module): + + def __init__(self, hidden_dim=128, input_dim=192 + 128): + super(SepConvGRU, self).__init__() + self.convz1 = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)) + self.convr1 = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)) + self.convq1 = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)) + + self.convz2 = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)) + self.convr2 = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)) + self.convq2 = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)) + + def forward(self, h, x): + # horizontal + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz1(hx)) + r = torch.sigmoid(self.convr1(hx)) + q = torch.tanh(self.convq1(torch.cat([r * h, x], dim=1))) + h = (1 - z) * h + z * q + + # vertical + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz2(hx)) + r = torch.sigmoid(self.convr2(hx)) + q = torch.tanh(self.convq2(torch.cat([r * h, x], dim=1))) + h = (1 - z) * h + z * q + + return h + + +class SmallMotionEncoder(nn.Module): + + def __init__(self, args): + super(SmallMotionEncoder, self).__init__() + cor_planes = args.corr_levels * (2 * args.corr_radius + 1)**2 + self.convc1 = nn.Conv2d(cor_planes, 96, 1, padding=0) + self.convf1 = nn.Conv2d(2, 64, 7, padding=3) + self.convf2 = nn.Conv2d(64, 32, 3, padding=1) + self.conv = nn.Conv2d(128, 80, 3, padding=1) + + def forward(self, flow, corr): + cor = F.relu(self.convc1(corr)) + flo = F.relu(self.convf1(flow)) + flo = F.relu(self.convf2(flo)) + cor_flo = torch.cat([cor, flo], dim=1) + out = F.relu(self.conv(cor_flo)) + return torch.cat([out, flow], dim=1) + + +class BasicMotionEncoder(nn.Module): + + def __init__(self, args): + super(BasicMotionEncoder, self).__init__() + cor_planes = args.corr_levels * (2 * args.corr_radius + 1)**2 + self.convc1 = nn.Conv2d(cor_planes, 256, 1, padding=0) + self.convc2 = nn.Conv2d(256, 192, 3, padding=1) + self.convf1 = nn.Conv2d(2, 128, 7, padding=3) + self.convf2 = nn.Conv2d(128, 64, 3, padding=1) + self.conv = nn.Conv2d(64 + 192, 128 - 2, 3, padding=1) + + def forward(self, flow, corr): + cor = F.relu(self.convc1(corr)) + cor = F.relu(self.convc2(cor)) + flo = F.relu(self.convf1(flow)) + flo = F.relu(self.convf2(flo)) + + cor_flo = torch.cat([cor, flo], dim=1) + out = F.relu(self.conv(cor_flo)) + return torch.cat([out, flow], dim=1) + + +class SmallUpdateBlock(nn.Module): + + def __init__(self, args, hidden_dim=96): + super(SmallUpdateBlock, self).__init__() + self.encoder = SmallMotionEncoder(args) + self.gru = ConvGRU(hidden_dim=hidden_dim, input_dim=82 + 64) + self.flow_head = FlowHead(hidden_dim, hidden_dim=128) + + def forward(self, net, inp, corr, flow): + motion_features = self.encoder(flow, corr) + inp = torch.cat([inp, motion_features], dim=1) + net = self.gru(net, inp) + delta_flow = self.flow_head(net) + + return net, None, delta_flow + + +class BasicUpdateBlock(nn.Module): + + def __init__(self, args, hidden_dim=128, input_dim=128): + super(BasicUpdateBlock, self).__init__() + self.args = args + self.encoder = BasicMotionEncoder(args) + self.gru = SepConvGRU( + hidden_dim=hidden_dim, input_dim=128 + hidden_dim) + self.flow_head = FlowHead(hidden_dim, hidden_dim=256) + + self.mask = nn.Sequential( + nn.Conv2d(128, 256, 3, padding=1), nn.ReLU(inplace=True), + nn.Conv2d(256, 64 * 9, 1, padding=0)) + + def forward(self, net, inp, corr, flow, upsample=True): + motion_features = self.encoder(flow, corr) + inp = torch.cat([inp, motion_features], dim=1) + + net = self.gru(net, inp) + delta_flow = self.flow_head(net) + + # scale mask to balence gradients + mask = .25 * self.mask(net) + return net, mask, delta_flow diff --git a/modelscope/models/cv/dense_optical_flow_estimation/core/utils/__init__.py b/modelscope/models/cv/dense_optical_flow_estimation/core/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/dense_optical_flow_estimation/core/utils/augmentor.py b/modelscope/models/cv/dense_optical_flow_estimation/core/utils/augmentor.py new file mode 100644 index 00000000..ff1b70dc --- /dev/null +++ b/modelscope/models/cv/dense_optical_flow_estimation/core/utils/augmentor.py @@ -0,0 +1,286 @@ +import math +import random + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image +from torchvision.transforms import ColorJitter + +cv2.setNumThreads(0) +cv2.ocl.setUseOpenCL(False) + + +class FlowAugmentor: + + def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=True): + + # spatial augmentation params + self.crop_size = crop_size + self.min_scale = min_scale + self.max_scale = max_scale + self.spatial_aug_prob = 0.8 + self.stretch_prob = 0.8 + self.max_stretch = 0.2 + + # flip augmentation params + self.do_flip = do_flip + self.h_flip_prob = 0.5 + self.v_flip_prob = 0.1 + + # photometric augmentation params + self.photo_aug = ColorJitter( + brightness=0.4, contrast=0.4, saturation=0.4, hue=0.5 / 3.14) + self.asymmetric_color_aug_prob = 0.2 + self.eraser_aug_prob = 0.5 + + def color_transform(self, img1, img2): + """ Photometric augmentation """ + + # asymmetric + if np.random.rand() < self.asymmetric_color_aug_prob: + img1 = np.array( + self.photo_aug(Image.fromarray(img1)), dtype=np.uint8) + img2 = np.array( + self.photo_aug(Image.fromarray(img2)), dtype=np.uint8) + + # symmetric + else: + image_stack = np.concatenate([img1, img2], axis=0) + image_stack = np.array( + self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8) + img1, img2 = np.split(image_stack, 2, axis=0) + + return img1, img2 + + def eraser_transform(self, img1, img2, bounds=[50, 100]): + """ Occlusion augmentation """ + + ht, wd = img1.shape[:2] + if np.random.rand() < self.eraser_aug_prob: + mean_color = np.mean(img2.reshape(-1, 3), axis=0) + for _ in range(np.random.randint(1, 3)): + x0 = np.random.randint(0, wd) + y0 = np.random.randint(0, ht) + dx = np.random.randint(bounds[0], bounds[1]) + dy = np.random.randint(bounds[0], bounds[1]) + img2[y0:y0 + dy, x0:x0 + dx, :] = mean_color + + return img1, img2 + + def spatial_transform(self, img1, img2, flow): + # randomly sample scale + ht, wd = img1.shape[:2] + min_scale = np.maximum((self.crop_size[0] + 8) / float(ht), + (self.crop_size[1] + 8) / float(wd)) + + scale = 2**np.random.uniform(self.min_scale, self.max_scale) + scale_x = scale + scale_y = scale + if np.random.rand() < self.stretch_prob: + scale_x *= 2**np.random.uniform(-self.max_stretch, + self.max_stretch) + scale_y *= 2**np.random.uniform(-self.max_stretch, + self.max_stretch) + + scale_x = np.clip(scale_x, min_scale, None) + scale_y = np.clip(scale_y, min_scale, None) + + if np.random.rand() < self.spatial_aug_prob: + # rescale the images + img1 = cv2.resize( + img1, + None, + fx=scale_x, + fy=scale_y, + interpolation=cv2.INTER_LINEAR) + img2 = cv2.resize( + img2, + None, + fx=scale_x, + fy=scale_y, + interpolation=cv2.INTER_LINEAR) + flow = cv2.resize( + flow, + None, + fx=scale_x, + fy=scale_y, + interpolation=cv2.INTER_LINEAR) + flow = flow * [scale_x, scale_y] + + if self.do_flip: + if np.random.rand() < self.h_flip_prob: # h-flip + img1 = img1[:, ::-1] + img2 = img2[:, ::-1] + flow = flow[:, ::-1] * [-1.0, 1.0] + + if np.random.rand() < self.v_flip_prob: # v-flip + img1 = img1[::-1, :] + img2 = img2[::-1, :] + flow = flow[::-1, :] * [1.0, -1.0] + + y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0]) + x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1]) + + img1 = img1[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + img2 = img2[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + flow = flow[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + + return img1, img2, flow + + def __call__(self, img1, img2, flow): + img1, img2 = self.color_transform(img1, img2) + img1, img2 = self.eraser_transform(img1, img2) + img1, img2, flow = self.spatial_transform(img1, img2, flow) + + img1 = np.ascontiguousarray(img1) + img2 = np.ascontiguousarray(img2) + flow = np.ascontiguousarray(flow) + + return img1, img2, flow + + +class SparseFlowAugmentor: + + def __init__(self, + crop_size, + min_scale=-0.2, + max_scale=0.5, + do_flip=False): + # spatial augmentation params + self.crop_size = crop_size + self.min_scale = min_scale + self.max_scale = max_scale + self.spatial_aug_prob = 0.8 + self.stretch_prob = 0.8 + self.max_stretch = 0.2 + + # flip augmentation params + self.do_flip = do_flip + self.h_flip_prob = 0.5 + self.v_flip_prob = 0.1 + + # photometric augmentation params + self.photo_aug = ColorJitter( + brightness=0.3, contrast=0.3, saturation=0.3, hue=0.3 / 3.14) + self.asymmetric_color_aug_prob = 0.2 + self.eraser_aug_prob = 0.5 + + def color_transform(self, img1, img2): + image_stack = np.concatenate([img1, img2], axis=0) + image_stack = np.array( + self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8) + img1, img2 = np.split(image_stack, 2, axis=0) + return img1, img2 + + def eraser_transform(self, img1, img2): + ht, wd = img1.shape[:2] + if np.random.rand() < self.eraser_aug_prob: + mean_color = np.mean(img2.reshape(-1, 3), axis=0) + for _ in range(np.random.randint(1, 3)): + x0 = np.random.randint(0, wd) + y0 = np.random.randint(0, ht) + dx = np.random.randint(50, 100) + dy = np.random.randint(50, 100) + img2[y0:y0 + dy, x0:x0 + dx, :] = mean_color + + return img1, img2 + + def resize_sparse_flow_map(self, flow, valid, fx=1.0, fy=1.0): + ht, wd = flow.shape[:2] + coords = np.meshgrid(np.arange(wd), np.arange(ht)) + coords = np.stack(coords, axis=-1) + + coords = coords.reshape(-1, 2).astype(np.float32) + flow = flow.reshape(-1, 2).astype(np.float32) + valid = valid.reshape(-1).astype(np.float32) + + coords0 = coords[valid >= 1] + flow0 = flow[valid >= 1] + + ht1 = int(round(ht * fy)) + wd1 = int(round(wd * fx)) + + coords1 = coords0 * [fx, fy] + flow1 = flow0 * [fx, fy] + + xx = np.round(coords1[:, 0]).astype(np.int32) + yy = np.round(coords1[:, 1]).astype(np.int32) + + v = (xx > 0) & (xx < wd1) & (yy > 0) & (yy < ht1) + xx = xx[v] + yy = yy[v] + flow1 = flow1[v] + + flow_img = np.zeros([ht1, wd1, 2], dtype=np.float32) + valid_img = np.zeros([ht1, wd1], dtype=np.int32) + + flow_img[yy, xx] = flow1 + valid_img[yy, xx] = 1 + + return flow_img, valid_img + + def spatial_transform(self, img1, img2, flow, valid): + # randomly sample scale + + ht, wd = img1.shape[:2] + min_scale = np.maximum((self.crop_size[0] + 1) / float(ht), + (self.crop_size[1] + 1) / float(wd)) + + scale = 2**np.random.uniform(self.min_scale, self.max_scale) + scale_x = np.clip(scale, min_scale, None) + scale_y = np.clip(scale, min_scale, None) + + if np.random.rand() < self.spatial_aug_prob: + # rescale the images + img1 = cv2.resize( + img1, + None, + fx=scale_x, + fy=scale_y, + interpolation=cv2.INTER_LINEAR) + img2 = cv2.resize( + img2, + None, + fx=scale_x, + fy=scale_y, + interpolation=cv2.INTER_LINEAR) + flow, valid = self.resize_sparse_flow_map( + flow, valid, fx=scale_x, fy=scale_y) + + if self.do_flip: + if np.random.rand() < 0.5: # h-flip + img1 = img1[:, ::-1] + img2 = img2[:, ::-1] + flow = flow[:, ::-1] * [-1.0, 1.0] + valid = valid[:, ::-1] + + margin_y = 20 + margin_x = 50 + + y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0] + margin_y) + x0 = np.random.randint(-margin_x, + img1.shape[1] - self.crop_size[1] + margin_x) + + y0 = np.clip(y0, 0, img1.shape[0] - self.crop_size[0]) + x0 = np.clip(x0, 0, img1.shape[1] - self.crop_size[1]) + + img1 = img1[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + img2 = img2[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + flow = flow[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + valid = valid[y0:y0 + self.crop_size[0], x0:x0 + self.crop_size[1]] + return img1, img2, flow, valid + + def __call__(self, img1, img2, flow, valid): + img1, img2 = self.color_transform(img1, img2) + img1, img2 = self.eraser_transform(img1, img2) + img1, img2, flow, valid = self.spatial_transform( + img1, img2, flow, valid) + + img1 = np.ascontiguousarray(img1) + img2 = np.ascontiguousarray(img2) + flow = np.ascontiguousarray(flow) + valid = np.ascontiguousarray(valid) + + return img1, img2, flow, valid diff --git a/modelscope/models/cv/dense_optical_flow_estimation/core/utils/flow_viz.py b/modelscope/models/cv/dense_optical_flow_estimation/core/utils/flow_viz.py new file mode 100644 index 00000000..46c92e34 --- /dev/null +++ b/modelscope/models/cv/dense_optical_flow_estimation/core/utils/flow_viz.py @@ -0,0 +1,132 @@ +# Flow visualization code used from https://github.com/tomrunia/OpticalFlow_Visualization + +# MIT License +# +# Copyright (c) 2018 Tom Runia +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to conditions. +# +# Author: Tom Runia +# Date Created: 2018-08-03 + +import numpy as np + + +def make_colorwheel(): + """ + Generates a color wheel for optical flow visualization as presented in: + Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007) + URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf + + Code follows the original C++ source code of Daniel Scharstein. + Code follows the the Matlab source code of Deqing Sun. + + Returns: + np.ndarray: Color wheel + """ + + RY = 15 + YG = 6 + GC = 4 + CB = 11 + BM = 13 + MR = 6 + + ncols = RY + YG + GC + CB + BM + MR + colorwheel = np.zeros((ncols, 3)) + col = 0 + + # RY + colorwheel[0:RY, 0] = 255 + colorwheel[0:RY, 1] = np.floor(255 * np.arange(0, RY) / RY) + col = col + RY + # YG + colorwheel[col:col + YG, 0] = 255 - np.floor(255 * np.arange(0, YG) / YG) + colorwheel[col:col + YG, 1] = 255 + col = col + YG + # GC + colorwheel[col:col + GC, 1] = 255 + colorwheel[col:col + GC, 2] = np.floor(255 * np.arange(0, GC) / GC) + col = col + GC + # CB + colorwheel[col:col + CB, 1] = 255 - np.floor(255 * np.arange(CB) / CB) + colorwheel[col:col + CB, 2] = 255 + col = col + CB + # BM + colorwheel[col:col + BM, 2] = 255 + colorwheel[col:col + BM, 0] = np.floor(255 * np.arange(0, BM) / BM) + col = col + BM + # MR + colorwheel[col:col + MR, 2] = 255 - np.floor(255 * np.arange(MR) / MR) + colorwheel[col:col + MR, 0] = 255 + return colorwheel + + +def flow_uv_to_colors(u, v, convert_to_bgr=False): + """ + Applies the flow color wheel to (possibly clipped) flow components u and v. + + According to the C++ source code of Daniel Scharstein + According to the Matlab source code of Deqing Sun + + Args: + u (np.ndarray): Input horizontal flow of shape [H,W] + v (np.ndarray): Input vertical flow of shape [H,W] + convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False. + + Returns: + np.ndarray: Flow visualization image of shape [H,W,3] + """ + flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8) + colorwheel = make_colorwheel() # shape [55x3] + ncols = colorwheel.shape[0] + rad = np.sqrt(np.square(u) + np.square(v)) + a = np.arctan2(-v, -u) / np.pi + fk = (a + 1) / 2 * (ncols - 1) + k0 = np.floor(fk).astype(np.int32) + k1 = k0 + 1 + k1[k1 == ncols] = 0 + f = fk - k0 + for i in range(colorwheel.shape[1]): + tmp = colorwheel[:, i] + col0 = tmp[k0] / 255.0 + col1 = tmp[k1] / 255.0 + col = (1 - f) * col0 + f * col1 + idx = (rad <= 1) + col[idx] = 1 - rad[idx] * (1 - col[idx]) + col[~idx] = col[~idx] * 0.75 # out of range + # Note the 2-i => BGR instead of RGB + ch_idx = 2 - i if convert_to_bgr else i + flow_image[:, :, ch_idx] = np.floor(255 * col) + return flow_image + + +def flow_to_image(flow_uv, clip_flow=None, convert_to_bgr=False): + """ + Expects a two dimensional flow image of shape. + + Args: + flow_uv (np.ndarray): Flow UV image of shape [H,W,2] + clip_flow (float, optional): Clip maximum of flow values. Defaults to None. + convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False. + + Returns: + np.ndarray: Flow visualization image of shape [H,W,3] + """ + assert flow_uv.ndim == 3, 'input flow must have three dimensions' + assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]' + if clip_flow is not None: + flow_uv = np.clip(flow_uv, 0, clip_flow) + u = flow_uv[:, :, 0] + v = flow_uv[:, :, 1] + rad = np.sqrt(np.square(u) + np.square(v)) + rad_max = np.max(rad) + epsilon = 1e-5 + u = u / (rad_max + epsilon) + v = v / (rad_max + epsilon) + return flow_uv_to_colors(u, v, convert_to_bgr) diff --git a/modelscope/models/cv/dense_optical_flow_estimation/core/utils/frame_utils.py b/modelscope/models/cv/dense_optical_flow_estimation/core/utils/frame_utils.py new file mode 100644 index 00000000..dac10fe1 --- /dev/null +++ b/modelscope/models/cv/dense_optical_flow_estimation/core/utils/frame_utils.py @@ -0,0 +1,142 @@ +import re +from os.path import * + +import cv2 +import numpy as np +from PIL import Image + +cv2.setNumThreads(0) +cv2.ocl.setUseOpenCL(False) + +TAG_CHAR = np.array([202021.25], np.float32) + + +def readFlow(fn): + """ Read .flo file in Middlebury format""" + # Code adapted from: + # http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy + + # WARNING: this will work on little-endian architectures (eg Intel x86) only! + # print 'fn = %s'%(fn) + with open(fn, 'rb') as f: + magic = np.fromfile(f, np.float32, count=1) + if 202021.25 != magic: + print('Magic number incorrect. Invalid .flo file') + return None + else: + w = np.fromfile(f, np.int32, count=1) + h = np.fromfile(f, np.int32, count=1) + # print 'Reading %d x %d flo file\n' % (w, h) + data = np.fromfile(f, np.float32, count=2 * int(w) * int(h)) + # Reshape data into 3D array (columns, rows, bands) + # The reshape here is for visualization, the original code is (w,h,2) + return np.resize(data, (int(h), int(w), 2)) + + +def readPFM(file): + file = open(file, 'rb') + + color = None + width = None + height = None + scale = None + endian = None + + header = file.readline().rstrip() + if header == b'PF': + color = True + elif header == b'Pf': + color = False + else: + raise Exception('Not a PFM file.') + + dim_match = re.match(rb'^(\d+)\s(\d+)\s$', file.readline()) + if dim_match: + width, height = map(int, dim_match.groups()) + else: + raise Exception('Malformed PFM header.') + + scale = float(file.readline().rstrip()) + if scale < 0: # little-endian + endian = '<' + scale = -scale + else: + endian = '>' # big-endian + + data = np.fromfile(file, endian + 'f') + shape = (height, width, 3) if color else (height, width) + + data = np.reshape(data, shape) + data = np.flipud(data) + return data + + +def writeFlow(filename, uv, v=None): + """ Write optical flow to file. + + If v is None, uv is assumed to contain both u and v channels, + stacked in depth. + Original code by Deqing Sun, adapted from Daniel Scharstein. + """ + nBands = 2 + + if v is None: + assert (uv.ndim == 3) + assert (uv.shape[2] == 2) + u = uv[:, :, 0] + v = uv[:, :, 1] + else: + u = uv + + assert (u.shape == v.shape) + height, width = u.shape + f = open(filename, 'wb') + # write the header + f.write(TAG_CHAR) + np.array(width).astype(np.int32).tofile(f) + np.array(height).astype(np.int32).tofile(f) + # arrange into matrix form + tmp = np.zeros((height, width * nBands)) + tmp[:, np.arange(width) * 2] = u + tmp[:, np.arange(width) * 2 + 1] = v + tmp.astype(np.float32).tofile(f) + f.close() + + +def readFlowKITTI(filename): + flow = cv2.imread(filename, cv2.IMREAD_ANYDEPTH | cv2.IMREAD_COLOR) + flow = flow[:, :, ::-1].astype(np.float32) + flow, valid = flow[:, :, :2], flow[:, :, 2] + flow = (flow - 2**15) / 64.0 + return flow, valid + + +def readDispKITTI(filename): + disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 256.0 + valid = disp > 0.0 + flow = np.stack([-disp, np.zeros_like(disp)], -1) + return flow, valid + + +def writeFlowKITTI(filename, uv): + uv = 64.0 * uv + 2**15 + valid = np.ones([uv.shape[0], uv.shape[1], 1]) + uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16) + cv2.imwrite(filename, uv[..., ::-1]) + + +def read_gen(file_name, pil=False): + ext = splitext(file_name)[-1] + if ext == '.png' or ext == '.jpeg' or ext == '.ppm' or ext == '.jpg': + return Image.open(file_name) + elif ext == '.bin' or ext == '.raw': + return np.load(file_name) + elif ext == '.flo': + return readFlow(file_name).astype(np.float32) + elif ext == '.pfm': + flow = readPFM(file_name).astype(np.float32) + if len(flow.shape) == 2: + return flow + else: + return flow[:, :, :-1] + return [] diff --git a/modelscope/models/cv/dense_optical_flow_estimation/core/utils/utils.py b/modelscope/models/cv/dense_optical_flow_estimation/core/utils/utils.py new file mode 100644 index 00000000..6228e6ef --- /dev/null +++ b/modelscope/models/cv/dense_optical_flow_estimation/core/utils/utils.py @@ -0,0 +1,93 @@ +import numpy as np +import torch +import torch.nn.functional as F +from scipy import interpolate + + +class InputPadder: + """ Pads images such that dimensions are divisible by 8 """ + + def __init__(self, dims, mode='sintel'): + self.ht, self.wd = dims[-2:] + pad_ht = (((self.ht // 8) + 1) * 8 - self.ht) % 8 + pad_wd = (((self.wd // 8) + 1) * 8 - self.wd) % 8 + if mode == 'sintel': + self._pad = [ + pad_wd // 2, pad_wd - pad_wd // 2, pad_ht // 2, + pad_ht - pad_ht // 2 + ] + else: + self._pad = [pad_wd // 2, pad_wd - pad_wd // 2, 0, pad_ht] + + def pad(self, *inputs): + return [F.pad(x, self._pad, mode='replicate') for x in inputs] + + def unpad(self, x): + ht, wd = x.shape[-2:] + c = [self._pad[2], ht - self._pad[3], self._pad[0], wd - self._pad[1]] + return x[..., c[0]:c[1], c[2]:c[3]] + + +def forward_interpolate(flow): + flow = flow.detach().cpu().numpy() + dx, dy = flow[0], flow[1] + + ht, wd = dx.shape + x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht)) + + x1 = x0 + dx + y1 = y0 + dy + + x1 = x1.reshape(-1) + y1 = y1.reshape(-1) + dx = dx.reshape(-1) + dy = dy.reshape(-1) + + valid = (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht) + x1 = x1[valid] + y1 = y1[valid] + dx = dx[valid] + dy = dy[valid] + + flow_x = interpolate.griddata((x1, y1), + dx, (x0, y0), + method='nearest', + fill_value=0) + + flow_y = interpolate.griddata((x1, y1), + dy, (x0, y0), + method='nearest', + fill_value=0) + + flow = np.stack([flow_x, flow_y], axis=0) + return torch.from_numpy(flow).float() + + +def bilinear_sampler(img, coords, mode='bilinear', mask=False): + """ Wrapper for grid_sample, uses pixel coordinates """ + H, W = img.shape[-2:] + xgrid, ygrid = coords.split([1, 1], dim=-1) + xgrid = 2 * xgrid / (W - 1) - 1 + ygrid = 2 * ygrid / (H - 1) - 1 + + grid = torch.cat([xgrid, ygrid], dim=-1) + img = F.grid_sample(img, grid, align_corners=True) + + if mask: + mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1) + return img, mask.float() + + return img + + +def coords_grid(batch, ht, wd, device): + coords = torch.meshgrid( + torch.arange(ht, device=device), torch.arange(wd, device=device)) + coords = torch.stack(coords[::-1], dim=0).float() + return coords[None].repeat(batch, 1, 1, 1) + + +def upflow8(flow, mode='bilinear'): + new_size = (8 * flow.shape[2], 8 * flow.shape[3]) + return 8 * F.interpolate( + flow, size=new_size, mode=mode, align_corners=True) diff --git a/modelscope/models/cv/dense_optical_flow_estimation/raft_model.py b/modelscope/models/cv/dense_optical_flow_estimation/raft_model.py new file mode 100644 index 00000000..2363092a --- /dev/null +++ b/modelscope/models/cv/dense_optical_flow_estimation/raft_model.py @@ -0,0 +1,52 @@ +import argparse +import os.path as osp + +import torch + +from modelscope.metainfo import Models +from modelscope.models.base.base_torch_model import TorchModel +from modelscope.models.builder import MODELS +from modelscope.models.cv.dense_optical_flow_estimation.core.raft import RAFT +from modelscope.outputs import OutputKeys +from modelscope.utils.constant import ModelFile, Tasks + + +@MODELS.register_module( + Tasks.dense_optical_flow_estimation, + module_name=Models.raft_dense_optical_flow_estimation) +class DenseOpticalFlowEstimation(TorchModel): + + def __init__(self, model_dir: str, **kwargs): + """str -- model file root.""" + super().__init__(model_dir, **kwargs) + + # build model + args = argparse.Namespace() + args.model = model_dir + args.small = False + args.mixed_precision = False + args.alternate_corr = False + self.model = torch.nn.DataParallel(RAFT(args)) + + model_path = osp.join(model_dir, ModelFile.TORCH_MODEL_FILE) + self.model.load_state_dict(torch.load(model_path)) + self.model = self.model.module + self.model.to('cuda') + self.model.eval() + + def forward(self, Inputs): + image1 = Inputs['image1'] + image2 = Inputs['image2'] + + flow_ups = self.model(image1, image2) + flow_up = flow_ups[-1] + + return flow_up + + def postprocess(self, inputs): + results = {OutputKeys.FLOWS: inputs} + return results + + def inference(self, data): + results = self.forward(data) + return results diff --git a/modelscope/outputs/outputs.py b/modelscope/outputs/outputs.py index 99569e06..4d3e0de3 100644 --- a/modelscope/outputs/outputs.py +++ b/modelscope/outputs/outputs.py @@ -25,6 +25,8 @@ class OutputKeys(object): MASKS = 'masks' DEPTHS = 'depths' DEPTHS_COLOR = 'depths_color' + FLOWS = 'flows' + FLOWS_COLOR = 'flows_color' NORMALS = 'normals' NORMALS_COLOR = 'normals_color' LAYOUT = 'layout' diff --git a/modelscope/pipelines/cv/dense_optical_flow_estimation_pipeline.py b/modelscope/pipelines/cv/dense_optical_flow_estimation_pipeline.py new file mode 100644 index 00000000..f734fd97 --- /dev/null +++ b/modelscope/pipelines/cv/dense_optical_flow_estimation_pipeline.py @@ -0,0 +1,147 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from typing import Any, Dict, Union + +import cv2 +import numpy as np +import PIL +import torch + +from modelscope.metainfo import Pipelines +from modelscope.outputs import OutputKeys +from modelscope.pipelines.base import Input, Model, Pipeline +from modelscope.pipelines.builder import PIPELINES +from modelscope.preprocessors import LoadImage +from modelscope.utils.constant import Tasks +from modelscope.utils.cv.image_utils import InputPadder, flow_to_color +from modelscope.utils.logger import get_logger + +logger = get_logger() + + +@PIPELINES.register_module( + Tasks.dense_optical_flow_estimation, + module_name=Pipelines.dense_optical_flow_estimation) +class DenseOpticalFlowEstimationPipeline(Pipeline): + r""" Card Detection Pipeline. + + Examples: + + >>> from modelscope.pipelines import pipeline + + >>> estimator = pipeline(Tasks.dense_optical_flow_estimation, model='Damo_XR_Lab/cv_raft_dense-optical-flow_things') + >>> estimator([[ + >>> 'modelscope/models/cv/dense_optical_flow_estimation/data/test/images/dense_flow1.png', + >>> 'modelscope/models/cv/dense_optical_flow_estimation/data/test/images/dense_flow2.png' + >>> ]]) + >>> [{'flows': tensor([[[[-1.6319, -1.6348, -1.6363, ..., -1.7191, -1.7136, -1.7085], + >>> [-1.6324, -1.6344, -1.6351, ..., -1.7110, -1.7048, -1.7005], + >>> [-1.6318, -1.6326, -1.6329, ..., -1.7080, -1.7050, -1.7031], + >>> ..., + >>> [-2.0998, -2.1007, -2.0958, ..., -1.4086, -1.4055, -1.3996], + >>> [-2.1043, -2.1031, -2.0988, ..., -1.4075, -1.4049, -1.3991], + >>> [-2.1016, -2.0985, -2.0939, ..., -1.4062, -1.4029, -1.3969]], + >>> + >>> [[ 0.0343, 0.0386, 0.0401, ..., 0.8053, 0.8050, 0.8057], + >>> [ 0.0311, 0.0354, 0.0369, ..., 0.8004, 0.8007, 0.8050], + >>> [ 0.0274, 0.0309, 0.0322, ..., 0.8007, 0.8016, 0.8080], + >>> ..., + >>> [ 0.5685, 0.5785, 0.5740, ..., 0.4003, 0.4153, 0.4365], + >>> [ 0.5994, 0.6000, 0.5899, ..., 0.4057, 0.4218, 0.4447], + >>> [ 0.6137, 0.6076, 0.5920, ..., 0.4147, 0.4299, 0.4538]]]], + >>> device='cuda:0'), 'flows_color': array([[[255, 249, 219], + >>> [255, 249, 219], + >>> [255, 249, 219], + >>> ..., + >>> [236, 255, 213], + >>> [236, 255, 213], + >>> [236, 255, 213]], + >>> + >>> [[255, 249, 219], + >>> [255, 249, 219], + >>> [255, 249, 219], + >>> ..., + >>> [236, 255, 213], + >>> [236, 255, 213], + >>> [236, 255, 213]], + >>> + >>> [[255, 249, 219], + >>> [255, 249, 219], + >>> [255, 249, 219], + >>> ..., + >>> [236, 255, 213], + >>> [236, 255, 213], + >>> [236, 255, 213]], + >>> + >>> ..., + >>> + >>> [[251, 255, 207], + >>> [251, 255, 207], + >>> [251, 255, 207], + >>> ..., + >>> [251, 255, 222], + >>> [251, 255, 222], + >>> [250, 255, 222]], + >>> + >>> [[250, 255, 207], + >>> [250, 255, 207], + >>> [250, 255, 207], + >>> ..., + >>> [251, 255, 222], + >>> [250, 255, 222], + >>> [249, 255, 222]], + >>> + >>> [[249, 255, 207], + >>> [249, 255, 207], + >>> [250, 255, 207], + >>> ..., + >>> [251, 255, 222], + >>> [250, 255, 222], + >>> [249, 255, 222]]], dtype=uint8)}] + """ + + def __init__(self, model: str, **kwargs): + """ + use `model` to create a image depth estimation pipeline for prediction + Args: + model: model id on modelscope hub. + """ + super().__init__(model=model, **kwargs) + + logger.info('dense optical flow estimation model, pipeline init') + + def load_image(self, img_name): + img = LoadImage.convert_to_ndarray(img_name).astype(np.float32) + img = img.transpose(2, 0, 1) + + return img + + def preprocess(self, input: Input) -> Dict[str, Any]: + img1 = self.load_image(input[0]) + img2 = self.load_image(input[1]) + + image1 = torch.from_numpy(img1)[None].cuda().float() + image2 = torch.from_numpy(img2)[None].cuda().float() + + padder = InputPadder(image1.shape) + image1, image2 = padder.pad(image1, image2) + + data = {'image1': image1, 'image2': image2} + + return data + + def forward(self, input: Dict[str, Any]) -> Dict[str, Any]: + flow_ups = self.model.inference(input) + results = flow_ups[-1] + + return results + + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + out = self.model.postprocess(inputs) + flows_color = flow_to_color([out[OutputKeys.FLOWS]]) + flows_color = flows_color[:, :, [2, 1, 0]] + outputs = { + OutputKeys.FLOWS: out[OutputKeys.FLOWS], + OutputKeys.FLOWS_COLOR: flows_color + } + + return outputs diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index 62a8dbd7..08850e5e 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -57,6 +57,7 @@ class CVTasks(object): semantic_segmentation = 'semantic-segmentation' image_driving_perception = 'image-driving-perception' image_depth_estimation = 'image-depth-estimation' + dense_optical_flow_estimation = 'dense-optical-flow-estimation' image_normal_estimation = 'image-normal-estimation' indoor_layout_estimation = 'indoor-layout-estimation' video_depth_estimation = 'video-depth-estimation' diff --git a/modelscope/utils/cv/image_utils.py b/modelscope/utils/cv/image_utils.py index 0efeae64..8eea4dea 100644 --- a/modelscope/utils/cv/image_utils.py +++ b/modelscope/utils/cv/image_utils.py @@ -7,6 +7,7 @@ import matplotlib import matplotlib.cm as cm import matplotlib.pyplot as plt import numpy as np +import torch.nn.functional as F from PIL import Image from modelscope.outputs import OutputKeys @@ -16,6 +17,30 @@ from modelscope.utils import logger as logging logger = logging.get_logger() +class InputPadder: + """ Pads images such that dimensions are divisible by 8 """ + + def __init__(self, dims, mode='sintel'): + self.ht, self.wd = dims[-2:] + pad_ht = (((self.ht // 8) + 1) * 8 - self.ht) % 8 + pad_wd = (((self.wd // 8) + 1) * 8 - self.wd) % 8 + if mode == 'sintel': + self._pad = [ + pad_wd // 2, pad_wd - pad_wd // 2, pad_ht // 2, + pad_ht - pad_ht // 2 + ] + else: + self._pad = [pad_wd // 2, pad_wd - pad_wd // 2, 0, pad_ht] + + def pad(self, *inputs): + return [F.pad(x, self._pad, mode='replicate') for x in inputs] + + def unpad(self, x): + ht, wd = x.shape[-2:] + c = [self._pad[2], ht - self._pad[3], self._pad[0], wd - self._pad[1]] + return x[..., c[0]:c[1], c[2]:c[3]] + + def numpy_to_cv2img(img_array): """to convert a np.array with shape(h, w) to cv2 img @@ -514,6 +539,127 @@ def depth_to_color(depth): return depth_color +def make_colorwheel(): + """ + Generates a color wheel for optical flow visualization as presented in: + Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007) + URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf + + Code follows the original C++ source code of Daniel Scharstein. + Code follows the the Matlab source code of Deqing Sun. + + Returns: + np.ndarray: Color wheel + """ + + RY = 15 + YG = 6 + GC = 4 + CB = 11 + BM = 13 + MR = 6 + + ncols = RY + YG + GC + CB + BM + MR + colorwheel = np.zeros((ncols, 3)) + col = 0 + + # RY + colorwheel[0:RY, 0] = 255 + colorwheel[0:RY, 1] = np.floor(255 * np.arange(0, RY) / RY) + col = col + RY + # YG + colorwheel[col:col + YG, 0] = 255 - np.floor(255 * np.arange(0, YG) / YG) + colorwheel[col:col + YG, 1] = 255 + col = col + YG + # GC + colorwheel[col:col + GC, 1] = 255 + colorwheel[col:col + GC, 2] = np.floor(255 * np.arange(0, GC) / GC) + col = col + GC + # CB + colorwheel[col:col + CB, 1] = 255 - np.floor(255 * np.arange(CB) / CB) + colorwheel[col:col + CB, 2] = 255 + col = col + CB + # BM + colorwheel[col:col + BM, 2] = 255 + colorwheel[col:col + BM, 0] = np.floor(255 * np.arange(0, BM) / BM) + col = col + BM + # MR + colorwheel[col:col + MR, 2] = 255 - np.floor(255 * np.arange(MR) / MR) + colorwheel[col:col + MR, 0] = 255 + return colorwheel + + +def flow_uv_to_colors(u, v, convert_to_bgr=False): + """ + Applies the flow color wheel to (possibly clipped) flow components u and v. + + According to the C++ source code of Daniel Scharstein + According to the Matlab source code of Deqing Sun + + Args: + u (np.ndarray): Input horizontal flow of shape [H,W] + v (np.ndarray): Input vertical flow of shape [H,W] + convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False. + + Returns: + np.ndarray: Flow visualization image of shape [H,W,3] + """ + flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8) + colorwheel = make_colorwheel() # shape [55x3] + ncols = colorwheel.shape[0] + rad = np.sqrt(np.square(u) + np.square(v)) + a = np.arctan2(-v, -u) / np.pi + fk = (a + 1) / 2 * (ncols - 1) + k0 = np.floor(fk).astype(np.int32) + k1 = k0 + 1 + k1[k1 == ncols] = 0 + f = fk - k0 + for i in range(colorwheel.shape[1]): + tmp = colorwheel[:, i] + col0 = tmp[k0] / 255.0 + col1 = tmp[k1] / 255.0 + col = (1 - f) * col0 + f * col1 + idx = (rad <= 1) + col[idx] = 1 - rad[idx] * (1 - col[idx]) + col[~idx] = col[~idx] * 0.75 # out of range + # Note the 2-i => BGR instead of RGB + ch_idx = 2 - i if convert_to_bgr else i + flow_image[:, :, ch_idx] = np.floor(255 * col) + return flow_image + + +def flow_to_image(flow_uv, clip_flow=None, convert_to_bgr=False): + """ + Expects a two dimensional flow image of shape. + + Args: + flow_uv (np.ndarray): Flow UV image of shape [H,W,2] + clip_flow (float, optional): Clip maximum of flow values. Defaults to None. + convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False. + + Returns: + np.ndarray: Flow visualization image of shape [H,W,3] + """ + assert flow_uv.ndim == 3, 'input flow must have three dimensions' + assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]' + if clip_flow is not None: + flow_uv = np.clip(flow_uv, 0, clip_flow) + u = flow_uv[:, :, 0] + v = flow_uv[:, :, 1] + rad = np.sqrt(np.square(u) + np.square(v)) + rad_max = np.max(rad) + epsilon = 1e-5 + u = u / (rad_max + epsilon) + v = v / (rad_max + epsilon) + return flow_uv_to_colors(u, v, convert_to_bgr) + + +def flow_to_color(flow): + flow = flow[0].permute(1, 2, 0).cpu().numpy() + flow_color = flow_to_image(flow) + return flow_color + + def show_video_depth_estimation_result(depths, video_save_path): height, width, layers = depths[0].shape out = cv2.VideoWriter(video_save_path, cv2.VideoWriter_fourcc(*'MP4V'), 25, diff --git a/modelscope/utils/pipeline_schema.json b/modelscope/utils/pipeline_schema.json index ace98cf9..c75fbfdf 100644 --- a/modelscope/utils/pipeline_schema.json +++ b/modelscope/utils/pipeline_schema.json @@ -1165,6 +1165,13 @@ "type": "object" } }, + "dense-optical-flow-estimation": { + "input": {}, + "parameters": {}, + "output": { + "type": "object" + } + }, "image-normal-estimation": { "input": {}, "parameters": {}, diff --git a/tests/pipelines/test_dense_optical_flow_estimation.py b/tests/pipelines/test_dense_optical_flow_estimation.py new file mode 100644 index 00000000..59ed8f12 --- /dev/null +++ b/tests/pipelines/test_dense_optical_flow_estimation.py @@ -0,0 +1,39 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import unittest + +import cv2 +import numpy as np + +from modelscope.outputs import OutputKeys +from modelscope.pipelines import pipeline +from modelscope.utils.constant import Tasks +from modelscope.utils.test_utils import test_level + + +class DenseOpticalFlowEstimationTest(unittest.TestCase): + + def setUp(self) -> None: + self.task = 'dense-optical-flow-estimation' + self.model_id = 'Damo_XR_Lab/cv_raft_dense-optical-flow_things' + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_dense_optical_flow_estimation(self): + input_location = [[ + 'data/test/images/dense_flow1.png', + 'data/test/images/dense_flow2.png', + # 'modelscope/models/cv/dense_optical_flow_estimation/data/test/images/dense_flow1.png', + # 'modelscope/models/cv/dense_optical_flow_estimation/data/test/images/dense_flow2.png' + ]] + estimator = pipeline( + Tasks.dense_optical_flow_estimation, model=self.model_id) + result = estimator(input_location) + # flow = result[0][OutputKeys.FLOWS] + flow_vis = result[0][OutputKeys.FLOWS_COLOR] + cv2.imwrite('result.jpg', flow_vis) + + print('test_dense_optical_flow_estimation DONE') + + +if __name__ == '__main__': + unittest.main()