Coverage for paintmaskanimator/auto_gpu_alignment.py: 0%

1"""線画をGPUに保持し、FFTの候補表を戻さず最良位置だけを返す。""" 

2from collections import OrderedDict 

3from typing import Any 

4 

5import numpy as np 

6 

7 

8class DeviceShift(tuple): 

9 device: Any 

10 gap: float 

11 def __new__(cls, dy, dx, device): 

12 result = super().__new__(cls, (int(dy), int(dx))) 

13 result.device = device 

14 return result 

15 

16 

17_PACK = r''' 

18extern "C" __global__ void pack( 

19 const unsigned char* target, const unsigned char* guide, const unsigned char* wide, 

20 int width, int height, int rows, int cols, int capacity, int job, 

21 int wx, int wy, int step, int tx0, int ty0, int tx1, int ty1, 

22 int h, int w, float weight, float* out) { 

23 int at=blockDim.x*blockIdx.x+threadIdx.x; if(at>=h*w) return; 

24 int y=at/w, x=at%w, sy=wy+y*step, sx=wx+x*step; 

25 int source=sy*width+sx, to=y*cols+x; 

26 out[job*rows*cols+to]=float(guide[source])*(1.0f-weight)+float(wide[source])*weight; 

27 out[(job+capacity)*rows*cols+to]= 

28 (sy>=ty0 && sy<ty1 && sx>=tx0 && sx<tx1) ? float(target[source]):0.0f; 

29} 

30''' 

31 

32 

33class ResidentAlignment: 

34 def __init__(self, cp): 

35 self.cp = cp 

36 self.frames = OrderedDict() 

37 self.pool = cp.cuda.MemoryPool() 

38 self.device = cp.cuda.Device().id 

39 self.pack_kernel = cp.RawKernel(_PACK, 'pack') 

40 self.uploads = 0 

41 self.returned_bytes = 0 

42 

43 def close(self): 

44 self.frames.clear() 

45 with self.cp.cuda.Device(self.device): 

46 self.pool.free_all_blocks() 

47 

48 def frame(self, cell): 

49 from .auto_alignment import boundary_key 

50 key = boundary_key(cell) 

51 if key in self.frames: 

52 result = self.frames.pop(key) 

53 self.frames[key] = result 

54 return result 

55 result = (self.cp.asarray(np.ascontiguousarray(cell.boundary)), 

56 self.cp.asarray(np.ascontiguousarray(cell.wide_boundary))) 

57 size = sum(a.nbytes for a in result) 

58 while self.frames and (len(self.frames) >= 12 or 

59 sum(sum(a.nbytes for a in item) for item in self.frames.values()) + size > 128 * 1024 * 1024): 

60 self.frames.popitem(last=False) 

61 self.frames[key] = result 

62 self.uploads += 1 

63 return result 

64 

65 def windows(self, bucket, batch, capacity): 

66 from . import auto_paint as ap 

67 cp = self.cp 

68 rows, cols, _limit = bucket 

69 with cp.cuda.Device(self.device), cp.cuda.using_allocator(self.pool.malloc): 

70 packed = cp.zeros((2 * capacity, rows, cols), dtype=cp.float32) 

71 for index, (_key, problem) in enumerate(batch): 

72 spec = problem.get('device_spec') 

73 h, w = problem['before'].shape 

74 if spec is None: 

75 packed[index, :h, :w] = cp.asarray(problem['after']) 

76 packed[index + capacity, :h, :w] = cp.asarray(problem['before']) 

77 continue 

78 target, guide, wx, wy, tx0, ty0, tx1, ty1 = spec 

79 first, _unused = self.frame(target) 

80 second, wide = self.frame(guide) 

81 height, width = target.values.shape 

82 values = [width, height, rows, cols, capacity, index, wx, wy, 

83 problem['step'], tx0, ty0, tx1, ty1, h, w] 

84 self.pack_kernel(((h*w+255)//256,), (256,), ( 

85 first, second, wide, *(np.int32(value) for value in values), 

86 np.float32(ap.NEAR_LINE_WEIGHT), packed, 

87 )) 

88 spectra = cp.fft.rfft2(packed, axes=(-2, -1)) 

89 correlation = cp.fft.irfft2( 

90 spectra[:capacity] * cp.conj(spectra[capacity:]), s=(rows, cols), axes=(-2, -1), 

91 ) 

92 ys = cp.asarray(np.array([problem['dys'] for _key, problem in batch])) % rows 

93 xs = cp.asarray(np.array([problem['dxs'] for _key, problem in batch])) % cols 

94 return correlation[cp.arange(len(batch))[:, None, None], ys[:, :, None], xs[:, None, :]] 

95 

96 def reduce(self, batch, windows): 

97 from . import auto_paint as ap 

98 cp = self.cp 

99 with cp.cuda.Device(self.device), cp.cuda.using_allocator(self.pool.malloc): 

100 # NumPyと同じfloat32の除算後、動きの罰則をfloat64で引く。 

101 total = cp.asarray(np.array([p['total'] for _key, p in batch], dtype=np.float32)) 

102 step = cp.asarray(np.array([p['step'] for _key, p in batch], dtype=np.int64)) 

103 dys = cp.asarray(np.array([p['dys'] for _key, p in batch], dtype=np.int64)) 

104 dxs = cp.asarray(np.array([p['dxs'] for _key, p in batch], dtype=np.int64)) 

105 scores = windows / total[:, None, None] - ap.MOTION_PENALTY * step[:, None, None] * ( 

106 cp.abs(dys)[:, :, None] + cp.abs(dxs)[:, None, :]) 

107 flat = scores.reshape(len(batch), -1) 

108 best = cp.argmax(flat, axis=1) 

109 at = cp.arange(len(batch)) 

110 peak = flat[at, best].copy() 

111 flat[at, best] = -cp.inf 

112 gap = peak - cp.max(flat, axis=1) 

113 y, x = best // windows.shape[2], best % windows.shape[2] 

114 result = cp.stack((dys[at, y] * step, dxs[at, x] * step, 

115 windows[at, y, x] / total, gap), axis=1) 

116 host = cp.asnumpy(result) 

117 self.returned_bytes += host.nbytes 

118 return host, result