Here’s the code that performs the classification task.
def classify_face(image):
device = torch.device("cpu")
img = process_image(image)
print('Image processed')
# img = image.unsqueeze_(0)
# img = image.float()
pred = model(img)[0]
# Apply NMS
pred = non_max_suppression(pred, 0.4, 0.5, classes = [0, 1, 2], agnostic = None )
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
#print(pred)
model.eval()
model.cpu()
print(pred)
# output = non_max_suppression(output, 0.4, 0.5, classes = class_names, agnostic = False)
#_, predicted = torch.max(output[0], 1)
#print(predicted.data[0], "predicted")
classification = torch.cat(pred)[:, -1]
index = int(classification)
print(names[index])
return names[index]
During prediction, pred consists of x1, y1, x2, y2, conf, and class.
Eg: pred = [tensor([[176.64380, 193.86154, 273.84702, 306.30405, 0.83492, 2.00000]])]
If there are no predictions made by the model then pred is simply empty.
Eg: pred = [tensor([], size=(0, 6))]
Presently my program stops prediction if it receives an empty tensor and throws an error:
Traceback (most recent call last):
File "WEBCAM_DETECT.py", line 168, in <module>
label = classify_face(frame)
File "WEBCAM_DETECT.py", line 150, in classify_face
index = int(classification)
ValueError: only one element tensors can be converted to Python scalars
How do I make my program sort of just ignore if there are no predictions made at a certain frame and continue onto the next frame?
My goal is to access the class from the tensor stored in pred, in integer.
Thank you for the help!