AvinaashA/DualView-BMI

1#General2import os3import numpy as np4import pandas as pd5 6#Feature extraction and Model7import torch8from facenet_pytorch import MTCNN, InceptionResnetV19from torchvision import transforms10from xgboost import XGBRegressor11 12#Image Processing and Display13from tqdm import tqdm14import warnings15from PIL import Image16import matplotlib.pyplot as plt17from skimage import io 18warnings.filterwarnings('ignore')19 20def load_and_process_image(image_path, device, mtcnn, resnet):21    try:22        img = Image.open(image_path)23        img_cropped = mtcnn(img)24        25        if img_cropped is None:26            print(f"No face detected in {image_path}")27            return None28            29        img_cropped = torch.unsqueeze(img_cropped, 0).to(device)30        with torch.no_grad():31            features = resnet(img_cropped)32        return features.cpu().numpy().flatten()33        34    except Exception as e:35        print(f"Error processing {image_path}: {str(e)}")36        return None37        38def extract_features(base_path, max_persons=1000):39    print("here")40    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')41    print(f"Using device: {device}")42    43    mtcnn = MTCNN(device=device)44    resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device)45    46    front_path = os.path.join(base_path, 'front')47    front_files = sorted(os.listdir(front_path))48 49    50    if max_persons: front_files = front_files[:max_persons]51    52    all_features = []53    processed_files = []54    55    for front_file in tqdm(front_files, desc="Processing images"):56        side_file = front_file  # Same filename in side folder57        58        front_features = load_and_process_image(59            os.path.join(front_path, front_file),60            device, mtcnn, resnet61        )62        #print("front-features")63        print(front_features)64        # print("Hi")65        66        side_features = load_and_process_image(67            os.path.join(base_path, 'side', side_file),68            device, mtcnn, resnet69        )70 71 72        print(side_features)73        if front_features is not None and side_features is not None:74            combined_features = np.concatenate([front_features, side_features])75            all_features.append(combined_features)76            processed_files.append(front_file)77 78 79    80    # Create feature column names81    front_cols = [f'front_feature_{i}' for i in range(512)]  # FaceNet outputs 512-D vectors82    side_cols = [f'side_feature_{i}' for i in range(512)]83    all_cols = front_cols + side_cols84    85    df = pd.DataFrame(all_features, columns=all_cols)86    df.insert(0, 'id', processed_files)87    88    return df89    90    91 92