Spaces:
Paused
Paused
| from PIL import Image | |
| import numpy as np | |
| import torch | |
| from cleanfid.features import build_feature_extractor, get_reference_statistics | |
| from cleanfid.fid import get_batch_features, fid_from_feats | |
| from cleanfid.resize import build_resizer | |
| """ | |
| A helper class that allowing adding the images one batch at a time. | |
| """ | |
| class CleanFID(): | |
| def __init__(self, mode="clean", model_name="inception_v3", device="cuda"): | |
| self.real_features = [] | |
| self.gen_features = [] | |
| self.mode = mode | |
| self.device = device | |
| if model_name=="inception_v3": | |
| self.feat_model = build_feature_extractor(mode, device) | |
| self.fn_resize = build_resizer(mode) | |
| elif model_name=="clip_vit_b_32": | |
| from cleanfid.clip_features import CLIP_fx, img_preprocess_clip | |
| clip_fx = CLIP_fx("ViT-B/32") | |
| self.feat_model = clip_fx | |
| self.fn_resize = img_preprocess_clip | |
| """ | |
| Funtion that takes an image (PIL.Image or np.array or torch.tensor) | |
| and returns the corresponding feature embedding vector. | |
| The image x is expected to be in range [0, 255] | |
| """ | |
| def compute_features(self, x): | |
| # if x is a PIL Image | |
| if isinstance(x, Image.Image): | |
| x_np = np.array(x) | |
| x_np_resized = self.fn_resize(x_np) | |
| x_t = torch.tensor(x_np_resized.transpose((2, 0, 1))).unsqueeze(0) | |
| x_feat = get_batch_features(x_t, self.feat_model, self.device) | |
| elif isinstance(x, np.ndarray): | |
| x_np_resized = self.fn_resize(x) | |
| x_t = torch.tensor(x_np_resized.transpose((2, 0, 1))).unsqueeze(0).to(self.device) | |
| # normalization happens inside the self.feat_model, expected image range here is [0,255] | |
| x_feat = get_batch_features(x_t, self.feat_model, self.device) | |
| elif isinstance(x, torch.Tensor): | |
| # pdb.set_trace() | |
| # add the batch dimension if x is passed in as C,H,W | |
| if len(x.shape)==3: | |
| x = x.unsqueeze(0) | |
| b,c,h,w = x.shape | |
| # convert back to np array and resize | |
| l_x_np_resized = [] | |
| for _ in range(b): | |
| x_np = x[_].cpu().numpy().transpose((1, 2, 0)) | |
| l_x_np_resized.append(self.fn_resize(x_np)[None,]) | |
| x_np_resized = np.concatenate(l_x_np_resized) | |
| x_t = torch.tensor(x_np_resized.transpose((0,3,1,2))).to(self.device) | |
| # normalization happens inside the self.feat_model, expected image range here is [0,255] | |
| x_feat = get_batch_features(x_t, self.feat_model, self.device) | |
| else: | |
| raise ValueError("image type could not be inferred") | |
| return x_feat | |
| """ | |
| Extract the faetures from x and add to the list of reference real images | |
| """ | |
| def add_real_images(self, x): | |
| x_feat = self.compute_features(x) | |
| self.real_features.append(x_feat) | |
| """ | |
| Extract the faetures from x and add to the list of generated images | |
| """ | |
| def add_gen_images(self, x): | |
| x_feat = self.compute_features(x) | |
| self.gen_features.append(x_feat) | |
| """ | |
| Compute FID between the real and generated images added so far | |
| """ | |
| def calculate_fid(self, verbose=True): | |
| feats1 = np.concatenate(self.real_features) | |
| feats2 = np.concatenate(self.gen_features) | |
| if verbose: | |
| print(f"# real images = {feats1.shape[0]}") | |
| print(f"# generated images = {feats2.shape[0]}") | |
| return fid_from_feats(feats1, feats2) | |
| """ | |
| Remove the real image features added so far | |
| """ | |
| def reset_real_features(self): | |
| self.real_features = [] | |
| """ | |
| Remove the generated image features added so far | |
| """ | |
| def reset_gen_features(self): | |
| self.gen_features = [] | |