-
Notifications
You must be signed in to change notification settings - Fork 6
/
Copy pathlightTest.py
108 lines (93 loc) · 2.95 KB
/
lightTest.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
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from torchvision import datasets,models,transforms
import torchvision
import os
import cv2
from PIL import Image
class Net(nn.Module):
def __init__(self):
super(Net,self).__init__()
# 1 input image channel 6 output channels 5*5 square convolution
#kernal
# 60*32*3->56*28*6
self.conv1=nn.Conv2d(3,6,5)
self.conv2=nn.Conv2d(6,16,5)
# an affine operation y=wx+b
self.fc1=nn.Linear(12*5*16,120)
self.fc2=nn.Linear(120,84)
self.fc3=nn.Linear(84,3)
def forward(self,x):
# Max pooling over a (2,2) window
# 56*28*6->28*14*6
x=F.max_pool2d(F.relu(self.conv1(x)),(2,2))
# 24*10*16->12*5*16
x=F.max_pool2d(F.relu(self.conv2(x)),2)
x=x.view(-1,16*5*12)
x=F.relu(self.fc1(x))
x=F.relu(self.fc2(x))
x=self.fc3(x)
x=F.log_softmax(x,dim=1)
return x
def num_flat_features(self,x):
size=x.size()[1:] # all dimensions except the batch dimension
num_features=1
for s in size:
num_features*=s
return num_features
image_size=(60,32)
process_transform=transforms.Compose([transforms.Resize(image_size),
transforms.ToTensor(),
transforms.Normalize([0.5,0.5,0.5], [0.5, 0.5, 0.5])])
classes=[
"red","green","yellow"
]
def detectLight(imgPath):
net=Net()
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
net.eval()
net.load_state_dict(torch.load("lightWeight.pth"))
img=Image.open(imgPath)
img_tensor=process_transform(img)
img_tensor.unsqueeze_(0)
img_tensor=img_tensor.to(device)
out=net(img_tensor)
_,predicted=torch.max(out,1)
#print(classes[predicted[0]])
def detectImg(img):
img=Image.fromarray(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
#img=Image.fromarray(img)
net=Net()
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
net.eval()
net.load_state_dict(torch.load("lightWeight.pth"))
img_tensor=process_transform(img)
img_tensor.unsqueeze_(0)
img_tensor=img_tensor.to(device)
out=net(img_tensor)
_,predicted=torch.max(out,1)
#print(classes[predicted[0]])
return classes[predicted[0]]
def detectImg2(img):
img=Image.fromarray(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
net=Net()
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
net.eval()
net.load_state_dict(torch.load("lightWeight_res.pth"))
img_tensor=process_transform(img)
img_tensor.unsqueeze_(0)
img_tensor=img_tensor.to(device)
out=net(img_tensor)
_,predicted=torch.max(out,1)
#print(classes[predicted[0]])
return classes[predicted[0]]
# detectLight("test03.jpg")
# img=cv2.imread("test01.jpg")
# detectImg(img)