import torch import mmdet import mmcv import mmengine import mmpretrain print("PyTorch version:", torch.__version__) print("GPU available:", torch.cuda.is_available()) print("GPU device count:", torch.cuda.device_count()) print("CUDA version:", torch.version.cuda) print("PyTorch CUDA archs:", torch.cuda.get_arch_list()) print("mmcv version:", mmcv.__version__) print("mmdet version:", mmdet.__version__) print("mmpretrain version:", mmpretrain.__version__) print("mmengine version:", mmengine.__version__) import os import gradio as gr import pandas as pd import numpy as np import cv2 import mmcv import json import pickle from mmdet.apis import DetInferencer from mmpretrain.apis import ImageClassificationInferencer from PIL import Image as IMG, ImageDraw, ImageFont config_file = "face_detection.py" checkpoint_file = "face_detection.pth" def face_detection_single_image(img): detector = DetInferencer(model=config_file, weights=checkpoint_file, device='cpu', show_progress = False) result = detector(mmcv.imread(img)) # img_array = result['visualization'][0] # # Convert the list into a NumPy array # img_array = np.array(img_array, dtype=np.uint8) # # Create an image from the array # output_image = IMG.fromarray(img_array) # print(result) bbox = [] # d = {} for i, r in enumerate(result['predictions'][0]['scores']): if r >= 0.5: bbox.append(result['predictions'][0]['bboxes'][i]) # Register mmpretrain modules mmpretrain.utils.register_all_modules() # Age and Gender age_inferencer = ImageClassificationInferencer( model='final_pretrained_eva02-base-p14_in1k.py', pretrained='eva02-base-p14_in1k-epoch_100.pth', device='cpu' ) gender_inferencer = ImageClassificationInferencer( model='convnext-v2-base_32xb32_in1k-384px.py', pretrained='convnext-v2-base_32xb32_in1k-384px-epoch12.pth', device='cpu' ) print("check-1") # original_image = IMG.fromarray(img) original_image = mmcv.imread(img) original_image = IMG.fromarray(original_image) age_pred = [] gender_pred = [] # bbox: [min_x, min_y, max_x, max_y] for b in bbox: cropped_image = np.array(original_image.crop((b[0], b[1], b[2], b[3]))) gender_results = gender_inferencer(cropped_image) gender_pred.append([gender_results[0]['pred_score'], gender_results[0]['pred_class']]) age_results = age_inferencer(cropped_image) age_pred.append([age_results[0]['pred_score'], age_results[0]['pred_class']]) print("check-2") # Convert the numpy array to a Pillow Image object image = IMG.fromarray(img) # Create a drawing object to draw on the image draw = ImageDraw.Draw(image) for i, box in enumerate(bbox): draw.rectangle(box, outline="green", width=3) # text_position_one = (box[0], box[1] - 10) # Slightly above the box # text_position_two = (box[0], box[1] - 20) # Slightly above the box gender_label = f"{'Female' if gender_pred[i][1] == 1 else 'Male'}: {round(gender_pred[i][0] * 100, 2)}%" age_label = f"Age: {age_pred[i][1]} ({round(age_pred[i][0] * 100, 2)}%)" text = f"{gender_label}\n{age_label}" font_size = max(2, (box[3] - box[1]) // 10) font = ImageFont.load_default() text_bbox = draw.textbbox((0, 0), text, font=font) text_width = text_bbox[2] - text_bbox[0] text_height = text_bbox[3] - text_bbox[1] text_position = (box[0], box[1] - text_height - 5) draw.rectangle([text_position[0], text_position[1], text_position[0] + text_width, text_position[1] + text_height], fill="black") # text_one = f"{round(age_pred[i][0], 2)},{age_pred[i][1]}" # text_two = f"{round(gender_pred[i][0], 2)},{gender_pred[i][1]}" draw.text(text_position, text, fill="white", font=font) # draw.text(text_position_one, text_one, fill="white") # draw.text(text_position_two, text_two, fill="white") print("check-3") return np.array(image) # def age_gender_single_image(custom_dict, gender_inferencer, age_inferencer, AGE_BIN_LABEL = ['0-11', '12-24', '25-29', '30-40', '41-64', '65+']): # ''' # Input1: dictionary containing annotations in the following format # { # 'img_path': { # FACE_NUMBER_IN_INT: { # 'face': { # 'coordinates': [x1, y1, x2, y2], # 'thr_score': CONFIDENCE_SCORE_IN_INT # } # } # } # } # Input2: Gender Inferencer in the following format # ImageClassificationInferencer(config_path, checkpoint_path, 'cuda') # Input3: Age Inferencer in the following format # ImageClassificationInferencer(config_path, checkpoint_path, 'cuda') # Optional: # AGE_BIN_LABEL: Can modify to match the number of age bins # Output1: dictionary with all the annotations and age, gender prediction # ''' # ### Apply age and gender inference on cropped_image # for key, val in custom_dict.items(): # selected_image = IMG.open(key) # if '.png' in key: # selected_image = selected_image.convert("RGB") # for k, v in val.items(): # # Testing prints # # Crop the image # cropped_image = np.array(selected_image.crop(( # v['face']['coordinates'][0], # v['face']['coordinates'][1], # v['face']['coordinates'][2], # v['face']['coordinates'][3]))) # # Gender Inference # gen = gender_inferencer(cropped_image) # if gen[0]['pred_label'] == 0: # is_female = 1 - gen[0]['pred_score'] # else: # is_female = gen[0]['pred_score'] # # data_dict[key][k]['gender'] = {} # custom_dict[key][k]['gender'] = { # 'is_female': is_female # } # # Age Inference # age = age_inferencer(cropped_image) # # data_dict[key][k]['age'] = {} # custom_dict[key][k]['age'] = pd.DataFrame( # { # 'age_bin': AGE_BIN_LABEL, # 'pred_score': age[0]['pred_scores'] # } # ) # return custom_dict def greet(name): return f"Hello {name}!" def main(): # if uploaded_image is not None: # config_file = "face_detection.py" # checkpoint_file = "face_detection.pth" # # inferencer = DetInferencer(model = config_file, weights = checkpoint_file, device = 'cpu') # tmp = face_detection_single_image(config_file, checkpoint_file, uploaded_image) # print(tmp) demo = gr.Interface( fn = face_detection_single_image, inputs = gr.Image(type="numpy"), outputs = gr.Image() # outputs = gr.Textbox() ) demo.launch(server_name="0.0.0.0", server_port=7860, debug=True) if __name__ == "__main__": main()