|
|
import matplotlib.pyplot as plt |
|
|
import numpy as np |
|
|
import torch |
|
|
from torchvision import transforms |
|
|
import sys |
|
|
import os |
|
|
import cv2 |
|
|
from FCN import FCN8s, FCN16s, FCN32s, FCNs, VGGNet |
|
|
from tqdm import tqdm |
|
|
import pandas as pd |
|
|
|
|
|
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
|
|
|
model = torch.load('./models/fcn_model_2.pt') |
|
|
model = model.to(device) |
|
|
transform = transforms.Compose([ |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) |
|
|
|
|
|
if __name__ =='__main__': |
|
|
occupations = ['backpacker', 'ballplayer', 'bartender', 'basketball_player', 'boatman', 'carpenter', 'cheerleader', 'climber', 'computer_user', 'craftsman', 'dancer', 'disk_jockey', 'doctor', 'drummer', 'electrician', 'farmer', 'fireman', 'flutist', 'gardener', 'guard', 'guitarist', 'gymnast', 'hairdresser', 'horseman', 'judge', 'laborer', 'lawman', 'lifeguard', 'machinist', 'motorcyclist', 'nurse', 'painter', 'patient', 'prayer', 'referee', 'repairman', 'reporter', 'retailer', 'runner', 'sculptor', 'seller', 'singer', 'skateboarder', 'soccer_player', 'soldier', 'speaker', 'student', 'teacher', 'tennis_player', 'trumpeter', 'waiter'] |
|
|
facet = pd.read_csv("../../../datasets/facet/annotations/annotations.csv", header=0).rename(columns={'Unnamed: 0': 'sample_idx'}) |
|
|
|
|
|
selected_categories = ['dancer', 'craftsman', 'gardener', 'retailer'] |
|
|
text_lists = [ |
|
|
[f"a photo of a {category}" for category in selected_categories], |
|
|
] |
|
|
|
|
|
root = "../../../datasets/facet/images_bb_small" |
|
|
|
|
|
|
|
|
for category in occupations: |
|
|
n_imgs = facet[facet['class1'] == category]['person_id'].shape[0] |
|
|
print(n_imgs) |
|
|
|
|
|
for id_ in tqdm(range(n_imgs)): |
|
|
img = facet[facet['class1'] == category].iloc[id_] |
|
|
if int(img['gender_presentation_masc']) == 1: |
|
|
gender = 'male' |
|
|
elif int(img['gender_presentation_fem']) == 1: |
|
|
gender = 'female' |
|
|
else: |
|
|
continue |
|
|
|
|
|
bb = eval(img["bounding_box"]) |
|
|
img_id = str(img['person_id']) |
|
|
|
|
|
img_path = os.path.join(root, img_id + ".jpg") |
|
|
|
|
|
|
|
|
|
|
|
save_file = './paper_output/FACET_FCN_skin' |
|
|
if not os.path.exists(save_file): |
|
|
os.makedirs(save_file) |
|
|
imgA = cv2.imread(img_path) |
|
|
print(img_id) |
|
|
size = imgA.shape |
|
|
|
|
|
imgA = cv2.resize(imgA, (160, 160)) |
|
|
|
|
|
imgA = transform(imgA) |
|
|
imgA = imgA.to(device) |
|
|
imgA = imgA.unsqueeze(0) |
|
|
output = model(imgA) |
|
|
output = torch.sigmoid(output) |
|
|
output_np = output.cpu().detach().numpy().copy() |
|
|
output_np = np.squeeze(output_np)*255 |
|
|
|
|
|
|
|
|
output_npA = output_np[0] |
|
|
output_npB = output_np[1] |
|
|
|
|
|
|
|
|
output_sigmoid = output_npA/(output_npA+output_npB)*255 |
|
|
output = cv2.resize(output_sigmoid,(size[1],size[0])) |
|
|
ret, output = cv2.threshold(output,127,255,cv2.THRESH_BINARY) |
|
|
|
|
|
cv2.imwrite(f"{save_file}/{img_id}.jpg", output) |
