Method forward

models/strand_codec.py:156–224

(self, gt_dict)

    def forward(self, gt_dict):

        points=gt_dict["strand_positions"]
        dirs=gt_dict["strand_directions"]

        #points
        points = points.permute(0, 2, 1) ## nr_strands, xyz, 100
        nr_strands = points.shape[0]
        #dirs
        last_dir = dirs[:, -1:, :]
        dirs = torch.cat([dirs, last_dir],1) # make the direction nr_strands, 100, 3
        dirs = dirs.permute(0, 2, 1)

        per_point_features = torch.cat([points, dirs] ,1)
        # per_point_features = points
        x=per_point_features

        # print("x",x.mean(),x.std())

        strand_features = self.cnn_encoder(x) # nr_strands, 128(nr_features), 3(elements per string)

        # print("strand_features after encoder", strand_features.mean(), strand_features.std())

        strand_features = strand_features.view(nr_strands, -1).contiguous()
        # strand_features = self.final_cnn_aggregator(strand_features) # outputs nr_strands x 128

        # print("strand_features after aggregate", strand_features.mean(), strand_features.std())

        strand_features_mean = self.aggregate_towards_mean(strand_features) # outputs nr_strands x 128
        s = self.pred_mean(strand_features_mean)
        
        # s = self.pred_mean(strand_features)
        # s=s

        # print("s mean and std", s.mean(), s.std())

        
        #pass the s through a tanh so it's bounded by -1,1 this makes it easier to consider it as an image later on when we train our diffusion model on scalp textures
        s=self.tanh(s)


        # exit(1)

        encoded_dict={}
        encoded_dict["z"]=s
        encoded_dict["z_no_eps"]=s

        if self.do_vae:
            s_mean = s
            # print("s_mean has mean std ", s_mean.mean(), s_mean.std())
            # s_logstd = 0.1 * self.pred_logstd(strand_features)
            strand_features_logstd = self.aggregate_towards_logstd(strand_features) # outputs nr_strands x 128
            s_logstd = -2.0 + 0.01*self.pred_logstd(strand_features_logstd) #start with logstd that is low so that initially the variance of the normal is also low
            encoded_dict["z_mean"] = s_mean
            encoded_dict["z_logstd"] = s_logstd
            # print("s_logstd has mean std ", s_logstd.mean(), s_logstd.std())
            if self.training:
                std = torch.exp(s_logstd)