-
Notifications
You must be signed in to change notification settings - Fork 1
/
script.py
114 lines (96 loc) · 6.47 KB
/
script.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
import torch
import torchvision.models as models
import torch.nn as nn
import sys, cv2
import numpy as np
from matplotlib import pyplot as plt
from PIL import Image
import torchvision.transforms as transforms
use_cuda = torch.cuda.is_available()
def image_to_tensor(img_path):
img = Image.open(img_path).convert('RGB')
transformations = transforms.Compose([transforms.Resize(size=224),
transforms.CenterCrop((224,224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])])
image_tensor = transformations(img)[:3,:,:].unsqueeze(0)
return image_tensor
def VGG16_predict(img_path):
image_tensor = image_to_tensor(img_path)
# move model inputs to cuda, if GPU available
if use_cuda:
image_tensor = image_tensor.cuda()
# get sample outputs
VGG16 = models.vgg16(pretrained=True)
if use_cuda:
VGG16 = VGG16.cuda()
output = VGG16(image_tensor)
# convert output probabilities to predicted class
_, preds_tensor = torch.max(output, 1)
pred = np.squeeze(preds_tensor.numpy()) if not use_cuda else np.squeeze(preds_tensor.cpu().numpy())
return int(pred) # predicted class index
### returns "True" if a dog is detected in the image stored at img_path
def dog_detector(img_path):
prediction = VGG16_predict(img_path)
return ((prediction >= 151) & (prediction <=268))
model_transfer = models.resnet50(pretrained=True)
use_cuda = torch.cuda.is_available()
if use_cuda:
model_transfer = model_transfer.cuda()
# Freeze parameters so we don't backprop through them
for param in model_transfer.parameters():
param.requires_grad = False
# Replace the last fully connected layer with a Linnear layer with 133 out features
model_transfer.fc = nn.Linear(2048, 133)
if use_cuda:
model_transfer = model_transfer.cuda()
# Load Pretrained Best Model
model_transfer.load_state_dict(torch.load('model_transfer.pt',map_location='cpu'))
class_names = ['Affenpinscher', 'Afghan hound', 'Airedale terrier', 'Akita', 'Alaskan malamute', 'American eskimo dog', 'American foxhound', 'American staffordshire terrier', 'American water spaniel', 'Anatolian shepherd dog', 'Australian cattle dog', 'Australian shepherd', 'Australian terrier', 'Basenji', 'Basset hound', 'Beagle', 'Bearded collie', 'Beauceron', 'Bedlington terrier', 'Belgian malinois', 'Belgian sheepdog', 'Belgian tervuren', 'Bernese mountain dog', 'Bichon frise', 'Black and tan coonhound', 'Black russian terrier', 'Bloodhound', 'Bluetick coonhound', 'Border collie', 'Border terrier', 'Borzoi', 'Boston terrier', 'Bouvier des flandres', 'Boxer', 'Boykin spaniel', 'Briard', 'Brittany', 'Brussels griffon', 'Bull terrier', 'Bulldog', 'Bullmastiff', 'Cairn terrier', 'Canaan dog', 'Cane corso', 'Cardigan welsh corgi', 'Cavalier king charles spaniel', 'Chesapeake bay retriever', 'Chihuahua', 'Chinese crested', 'Chinese shar-pei', 'Chow chow', 'Clumber spaniel', 'Cocker spaniel', 'Collie', 'Curly-coated retriever', 'Dachshund', 'Dalmatian', 'Dandie dinmont terrier', 'Doberman pinscher', 'Dogue de bordeaux', 'English cocker spaniel', 'English setter', 'English springer spaniel', 'English toy spaniel', 'Entlebucher mountain dog', 'Field spaniel', 'Finnish spitz', 'Flat-coated retriever', 'French bulldog', 'German pinscher', 'German shepherd dog', 'German shorthaired pointer', 'German wirehaired pointer', 'Giant schnauzer', 'Glen of imaal terrier', 'Golden retriever', 'Gordon setter', 'Great dane', 'Great pyrenees', 'Greater swiss mountain dog', 'Greyhound', 'Havanese', 'Ibizan hound', 'Icelandic sheepdog', 'Irish red and white setter', 'Irish setter', 'Irish terrier', 'Irish water spaniel', 'Irish wolfhound', 'Italian greyhound', 'Japanese chin', 'Keeshond', 'Kerry blue terrier', 'Komondor', 'Kuvasz', 'Labrador retriever', 'Lakeland terrier', 'Leonberger', 'Lhasa apso', 'Lowchen', 'Maltese', 'Manchester terrier', 'Mastiff', 'Miniature schnauzer', 'Neapolitan mastiff', 'Newfoundland', 'Norfolk terrier', 'Norwegian buhund', 'Norwegian elkhound', 'Norwegian lundehund', 'Norwich terrier', 'Nova scotia duck tolling retriever', 'Old english sheepdog', 'Otterhound', 'Papillon', 'Parson russell terrier', 'Pekingese', 'Pembroke welsh corgi', 'Petit basset griffon vendeen', 'Pharaoh hound', 'Plott', 'Pointer', 'Pomeranian', 'Poodle', 'Portuguese water dog', 'Saint bernard', 'Silky terrier', 'Smooth fox terrier', 'Tibetan mastiff', 'Welsh springer spaniel', 'Wirehaired pointing griffon', 'Xoloitzcuintli', 'Yorkshire terrier']
def predict_breed_transfer(img_path):
# load the image and return the predicted breed
image_tensor = image_to_tensor(img_path)
if use_cuda:
image_tensor = image_tensor.cuda()
output = model_transfer(image_tensor)
_, preds_tensor = torch.max(output, 1)
pred = np.squeeze(preds_tensor.numpy()) if not use_cuda else np.squeeze(preds_tensor.cpu().numpy())
return class_names[pred]
def display_image(img_path, title="Title"):
image = Image.open(img_path)
plt.title(title)
plt.imshow(image)
plt.show()
def haar_face_detector(img_path):
haar_face_cascade = cv2.CascadeClassifier('haarcascades/haarcascade_frontalface_alt.xml')
img = cv2.imread(img_path)
gray = convertToRGB(img)
faces = haar_face_cascade.detectMultiScale(gray)
return len(faces) > 0
def convertToRGB(img):
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
### TODO: Write your algorithm.
### Feel free to use as many code cells as needed.
def run_app(img_path):
## handle cases for a human face, dog, and neither
# check if image has juman faces:
print("IMAGE PATH: ", img_path)
if (haar_face_detector(img_path)):
print("This is human!")
predicted_breed = predict_breed_transfer(img_path)
#display_image(img_path, title=f"Predicted:{predicted_breed}")
print("You look like a ...", end=" ")
print(predicted_breed.upper())
# check if image has dogs:
elif dog_detector(img_path):
print("This is dog!")
predicted_breed = predict_breed_transfer(img_path)
#display_image(img_path, title=f"Predicted:{predicted_breed}")
print("Your breed is most likely ...", end=" ")
print(predicted_breed.upper())
else:
print("Sorry. Something went wrong..")
display_image(img_path, title="...")
print("Try something else!")
run_app(sys.argv[1])