1/*
2 *  Copyright 2017 The WebRTC project authors. All Rights Reserved.
3 *
4 *  Use of this source code is governed by a BSD-style license
5 *  that can be found in the LICENSE file in the root of the source
6 *  tree. An additional intellectual property rights grant can be found
7 *  in the file PATENTS.  All contributing project authors may
8 *  be found in the AUTHORS file in the root of the source tree.
9 */
10
11#import "RTCMTLI420Renderer.h"
12
13#import <Metal/Metal.h>
14#import <MetalKit/MetalKit.h>
15
16#import "base/RTCI420Buffer.h"
17#import "base/RTCLogging.h"
18#import "base/RTCVideoFrame.h"
19#import "base/RTCVideoFrameBuffer.h"
20
21#import "RTCMTLRenderer+Private.h"
22
23static NSString *const shaderSource = MTL_STRINGIFY(
24    using namespace metal;
25
26    typedef struct {
27      packed_float2 position;
28      packed_float2 texcoord;
29    } Vertex;
30
31    typedef struct {
32      float4 position[[position]];
33      float2 texcoord;
34    } Varyings;
35
36    vertex Varyings vertexPassthrough(constant Vertex *verticies[[buffer(0)]],
37                                      unsigned int vid[[vertex_id]]) {
38      Varyings out;
39      constant Vertex &v = verticies[vid];
40      out.position = float4(float2(v.position), 0.0, 1.0);
41      out.texcoord = v.texcoord;
42
43      return out;
44    }
45
46    fragment half4 fragmentColorConversion(
47        Varyings in[[stage_in]],
48        texture2d<float, access::sample> textureY[[texture(0)]],
49        texture2d<float, access::sample> textureU[[texture(1)]],
50        texture2d<float, access::sample> textureV[[texture(2)]]) {
51      constexpr sampler s(address::clamp_to_edge, filter::linear);
52      float y;
53      float u;
54      float v;
55      float r;
56      float g;
57      float b;
58      // Conversion for YUV to rgb from http://www.fourcc.org/fccyvrgb.php
59      y = textureY.sample(s, in.texcoord).r;
60      u = textureU.sample(s, in.texcoord).r;
61      v = textureV.sample(s, in.texcoord).r;
62      u = u - 0.5;
63      v = v - 0.5;
64      r = y + 1.403 * v;
65      g = y - 0.344 * u - 0.714 * v;
66      b = y + 1.770 * u;
67
68      float4 out = float4(r, g, b, 1.0);
69
70      return half4(out);
71    });
72
73@implementation RTCMTLI420Renderer {
74  // Textures.
75  id<MTLTexture> _yTexture;
76  id<MTLTexture> _uTexture;
77  id<MTLTexture> _vTexture;
78
79  MTLTextureDescriptor *_descriptor;
80  MTLTextureDescriptor *_chromaDescriptor;
81
82  int _width;
83  int _height;
84  int _chromaWidth;
85  int _chromaHeight;
86}
87
88#pragma mark - Virtual
89
90- (NSString *)shaderSource {
91  return shaderSource;
92}
93
94- (void)getWidth:(nonnull int *)width
95          height:(nonnull int *)height
96       cropWidth:(nonnull int *)cropWidth
97      cropHeight:(nonnull int *)cropHeight
98           cropX:(nonnull int *)cropX
99           cropY:(nonnull int *)cropY
100         ofFrame:(nonnull RTC_OBJC_TYPE(RTCVideoFrame) *)frame {
101  *width = frame.width;
102  *height = frame.height;
103  *cropWidth = frame.width;
104  *cropHeight = frame.height;
105  *cropX = 0;
106  *cropY = 0;
107}
108
109- (BOOL)setupTexturesForFrame:(nonnull RTC_OBJC_TYPE(RTCVideoFrame) *)frame {
110  if (![super setupTexturesForFrame:frame]) {
111    return NO;
112  }
113
114  id<MTLDevice> device = [self currentMetalDevice];
115  if (!device) {
116    return NO;
117  }
118
119  id<RTC_OBJC_TYPE(RTCI420Buffer)> buffer = [frame.buffer toI420];
120
121  // Luma (y) texture.
122  if (!_descriptor || _width != frame.width || _height != frame.height) {
123    _width = frame.width;
124    _height = frame.height;
125    _descriptor = [MTLTextureDescriptor texture2DDescriptorWithPixelFormat:MTLPixelFormatR8Unorm
126                                                                     width:_width
127                                                                    height:_height
128                                                                 mipmapped:NO];
129    _descriptor.usage = MTLTextureUsageShaderRead;
130    _yTexture = [device newTextureWithDescriptor:_descriptor];
131  }
132
133  // Chroma (u,v) textures
134  [_yTexture replaceRegion:MTLRegionMake2D(0, 0, _width, _height)
135               mipmapLevel:0
136                 withBytes:buffer.dataY
137               bytesPerRow:buffer.strideY];
138
139  if (!_chromaDescriptor || _chromaWidth != frame.width / 2 || _chromaHeight != frame.height / 2) {
140    _chromaWidth = frame.width / 2;
141    _chromaHeight = frame.height / 2;
142    _chromaDescriptor =
143        [MTLTextureDescriptor texture2DDescriptorWithPixelFormat:MTLPixelFormatR8Unorm
144                                                           width:_chromaWidth
145                                                          height:_chromaHeight
146                                                       mipmapped:NO];
147    _chromaDescriptor.usage = MTLTextureUsageShaderRead;
148    _uTexture = [device newTextureWithDescriptor:_chromaDescriptor];
149    _vTexture = [device newTextureWithDescriptor:_chromaDescriptor];
150  }
151
152  [_uTexture replaceRegion:MTLRegionMake2D(0, 0, _chromaWidth, _chromaHeight)
153               mipmapLevel:0
154                 withBytes:buffer.dataU
155               bytesPerRow:buffer.strideU];
156  [_vTexture replaceRegion:MTLRegionMake2D(0, 0, _chromaWidth, _chromaHeight)
157               mipmapLevel:0
158                 withBytes:buffer.dataV
159               bytesPerRow:buffer.strideV];
160
161  return (_uTexture != nil) && (_yTexture != nil) && (_vTexture != nil);
162}
163
164- (void)uploadTexturesToRenderEncoder:(id<MTLRenderCommandEncoder>)renderEncoder {
165  [renderEncoder setFragmentTexture:_yTexture atIndex:0];
166  [renderEncoder setFragmentTexture:_uTexture atIndex:1];
167  [renderEncoder setFragmentTexture:_vTexture atIndex:2];
168}
169
170@end
171