MCPcopy Create free account
hub / github.com/MotrixLab/ViMoGen / get_bucket_id

Method get_bucket_id

datasets/bucket.py:81–128  ·  view source on GitHub ↗
(self, T, H, W, frame_interval=1, seed=None)

Source from the content-addressed store, hash-verified

79 get_logger().info('Number of buckets: %s', num_bucket)
80
81 def get_bucket_id(self, T, H, W, frame_interval=1, seed=None):
82 resolution = H * W
83 approx = 0.8
84
85 fail = True
86 for hw_id, t_criteria in self.bucket_probs.items():
87 if resolution < self.hw_criteria[hw_id] * approx:
88 continue
89
90 # if sample is an image
91 if T == 1:
92 if 1 in t_criteria:
93 rng = np.random.default_rng(seed +
94 self.bucket_id[hw_id][1])
95 if rng.random() < t_criteria[1]:
96 fail = False
97 t_id = 1
98 break
99 else:
100 continue
101
102 # otherwise, find suitable t_id for video
103 t_fail = True
104 for t_id, prob in t_criteria.items():
105 rng = np.random.default_rng(seed + self.bucket_id[hw_id][t_id])
106 if isinstance(prob, tuple):
107 prob_t = prob[1]
108 if rng.random() > prob_t:
109 continue
110 if T > t_id * frame_interval and t_id != 1:
111 t_fail = False
112 break
113 if t_fail:
114 continue
115
116 # leave the loop if prob is high enough
117 if isinstance(prob, tuple):
118 prob = prob[0]
119 if prob >= 1 or rng.random() < prob:
120 fail = False
121 break
122 if fail:
123 return None
124
125 # get aspect ratio id
126 ar_criteria = self.ar_criteria[hw_id][t_id]
127 ar_id = get_closest_ratio(H, W, ar_criteria)
128 return hw_id, t_id, ar_id
129
130 def get_thw(self, bucket_id):
131 assert len(bucket_id) == 3

Callers

nothing calls this directly

Calls 1

get_closest_ratioFunction · 0.85

Tested by

no test coverage detected