Caffe/classify.py at master · PaperAnalyticalDeviceND/Caffe · GitHub

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
import os
import numpy as np
import matplotlib.pyplot as plt
import math
import copy
from PIL import Image

#setup drugs list
drugs = ["Acetaminophen",
         "AcetylsalicylicAcid",
         "Amodiaquine",
         "Amoxicillin",
         "Ampicillin",
         "Artesunate",
         "CalciumCarbonate",
         "CornStarch",
         "Diethylcarbamazine",
         "Ethambutol",
         "Isoniazid",
         "Rifampicin",
         "Tetracycline",
         "Azithromycin",
         "Chloramphenicol",
         "Chloroquine",
         "Ciprofloxacin",
         "DIWater",
         "DriedWheatStarch",
         "PenicillinG",
         "PotatoStarch",
         "Primaquine",
         "Quinine",
         "Streptomycin",
         "Sulfadoxine",
         "Talc"
         ]

#caffe_root = '/home/jsweet/Documents/Code/caffe/'  # this file is expected to be in {caffe_root}/examples
import sys
#sys.path.insert(0, caffe_root + 'python')
#set to my folder
os.chdir(os.path.dirname(os.path.abspath(__file__)))

import caffe

#get processed image
img = Image.open(sys.argv[1])

#crop comparison
img = img.crop((72, 359+5, 72+636, 359+5+490))
img.save('tmp/test.png')

caffe.set_mode_cpu()
net = caffe.Classifier('deploy.prototxt', 'Sandipan1_Full_26Drugs_iter_50000.caffemodel',
                       mean=np.load('imagenet_mean.npy').mean(1).mean(1),
                       channel_swap=(2,1,0),
                       raw_scale=255,
                       image_dims=(256, 256))


input_image = caffe.io.load_image('tmp/test.png')
prediction = net.predict([input_image])

temppred = copy.deepcopy(prediction[0])
pClass1 = temppred.argmax()

temppred[pClass1] = 0
pClass2 = temppred.argmax()

temppred[pClass2] = 0
pClass3 = temppred.argmax()

print file
print '\tClass:', pClass1, pClass2, pClass3
print '\tProbability', prediction[0][pClass1], prediction[0][pClass2], prediction[0][pClass3]
print '\tDrug', drugs[pClass1], drugs[pClass2], drugs[pClass3]