Commit a9d2edab authored by s8_filimo_m's avatar s8_filimo_m
Browse files

Upload New File

parent fa821dbd
Pipeline #1229 canceled with stages
import numpy as np
import cv2
import imutils
import argparse
from models import *
from torchvision import models, transforms
from torch.autograd import Variable
from PIL import Image
import torch.nn as nn
outputFile = cv2.VideoWriter('appsrc ! videoconvert ! x264enc bitrate=10000 ! mp4mux ! filesink location=/', cv2.CAP_GSTREAMER, 30, (352, 608))
loader = transforms.Compose([transforms.Resize(264), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
def image_loader(image_data):
image_data = np.flip(image_data, axis=2)
img = Image.fromarray(image_data, 'RGB')
img = loader(img).float()
img = Variable(img, requires_grad=True)
img = img.unsqueeze(0)
return img.cuda()
def detect(videoFile, showVideo = 0):
cap = cv2.VideoCapture(videoFile)
img_size = (608, 352)
weights = "/home/Documents/Maksims/CourseWork/handGestures/FinalVersion/yolov3/weights/best.pt"
config = "/home/Documents/Maksims/CourseWork/handGestures/FinalVersion/yolov3/cfg/yolov3-spp-1cls-608.cfg"
# Initialize
device = torch_utils.select_device(device='')
# Initialize model
model = Darknet(config, img_size)
#Class names
thisdict = {
"0": "A",
"1": "B",
"2": "C",
"3": "D",
"4": "E"
}
# Load weights
attempt_download(weights)
if weights.endswith('.pt'): # pytorch format
model.load_state_dict(torch.load(weights, map_location=device)['model'])
else: # darknet format
load_darknet_weights(model, weights)
# Load weights for classifier
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
class_names = ['A', 'B', 'C', 'D', 'E']
model_ft = models.resnet34(pretrained=True)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, len(class_names))
model_ft = model_ft.to(device)
model_ft.load_state_dict(torch.load('/home/Documents/Maksims/CourseWork/handGestures/FinalVersion/yolov3/weights/handClassifier.pt'))
model_ft.eval()
# Eval mode
model.to(device).eval()
# Run inference
while(cap.isOpened()):
ret, frame = cap.read()
if np.shape(frame) == ():
break
frame = cv2.resize(frame,(img_size[0], img_size[1]))
frame = imutils.rotate_bound(frame, 90)
ogFrame = frame
frame = frame.transpose((2,0,1))
img = torch.from_numpy(frame).float().div(255)
img = img.unsqueeze(0)
img = img.to(device)
# Inference
pred = model(img)[0]
# Apply NMS Conf IoU
pred = non_max_suppression(pred, 0.3, 0.6, classes='', agnostic=False)
# Process detections
total = 0
start = time.perf_counter()
for i, det in enumerate(pred): # detections per image
if det is not None and len(det):
# Write results
for *xyxy, conf, cls in det:
confCPU = str(conf.cpu().numpy())
c1, c2, c3, c4 = int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])
cutout = ogFrame[c2:c4, c1:c3]
if 0 in cutout.shape:
pass
else:
img = image_loader(cutout)
outputs = model_ft(img)
outputsCpu = outputs.data.cpu().numpy()[0]
testarg = np.argmax(outputsCpu)
telabel = thisdict[str(testarg)]
confCPU = confCPU[:4]
labelVal = telabel + ": " + confCPU
plot_one_box(xyxy, ogFrame, label=labelVal, color=[255, 51, 204])
outputFile.write(ogFrame)
if showVideo == 1:
cv2.imshow('frame',ogFrame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
outputFile.release()
cv2.destroyAllWindows()
if __name__ == '__main__':
with torch.no_grad():
detect(videoFile)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment