src = test_dataset[i][0]
img_name = os.path.basename(test_dataset.ids[i][0]).split('.')[0]
image = preproc_for_test(src, opt.min_size, opt.mean)
image = torch.from_numpy(image).to(device)
with torch.no_grad():
loc, conf = model(image.unsqueeze(0))
loc = loc[0]
conf = conf[0]
conf = F.softmax(conf, dim=1)
conf = conf.cpu().numpy()
loc = loc.cpu().numpy()
decode_loc = multibox_encoder.decode(loc)
gt_boxes, gt_confs, gt_labels = detect(decode_loc,
conf,
nms_threshold=0.5,
gt_threshold=0.01)
#no object detected
if len(gt_boxes) == 0:
continue
h, w = src.shape[:2]
gt_boxes[:, 0] = gt_boxes[:, 0] * w
gt_boxes[:, 1] = gt_boxes[:, 1] * h
gt_boxes[:, 2] = gt_boxes[:, 2] * w
gt_boxes[:, 3] = gt_boxes[:, 3] * h
for box, label, score in zip(gt_boxes, gt_labels, gt_confs):
print(img_name,
"{:.3f}".format(score),
def utilspytorch(lines,classes=None):
#classes = ['forklift','person','pallet','box']
newpath = './output'
# delete any temp directory and its contents
if os.path.exists(newpath):
shutil.rmtree('./output')
# if directory is not yet created, then create one
if not os.path.exists(newpath):
os.makedirs(newpath)
ct = CentroidTracker()
cap = cv2.VideoCapture('temp/video.mp4')
time.sleep(2.0)
count = 0
total_points = {}
color_id = {}
while True:
# read the next frame from the video stream and resize it
# frame = vs.read()
ret, src = cap.read()
if not ret:
#do a bit of cleanup
break
cv2.destroyAllWindows()
#src = cv2.imread(image,cv2.IMREAD_COLOR)
image = preproc_for_test(src, opt.min_size, opt.mean)
image = torch.from_numpy(image)
with torch.no_grad():
loc, conf = model(image.unsqueeze(0))
loc = loc[0]
conf = conf[0]
conf = F.softmax(conf, dim=1)
conf = conf.numpy()
loc = loc.numpy()
decode_loc = multibox_encoder.decode(loc)
gt_boxes, gt_confs, gt_labels = detect(decode_loc, conf, nms_threshold=0.45, gt_threshold=0.1)
h, w = 320, 576
src = cv2.resize(src, (576, 320))
gt_boxes[:, 0] = gt_boxes[:, 0] * w
gt_boxes[:, 1] = gt_boxes[:, 1] * h
gt_boxes[:, 2] = gt_boxes[:, 2] * w
gt_boxes[:, 3] = gt_boxes[:, 3] * h
gt_boxes = gt_boxes.astype(int)
rects = {}
rects_counter = 0
for box, label, score in zip(gt_boxes, gt_labels, gt_confs):
if VOC_LABELS[label] in classes:
object_type = VOC_LABELS[label]
if score > 0.3:
image = cv2.rectangle(src, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 2)
rect = np.array([box[0], box[1], box[2], box[3]])
rects[rects_counter] = (rect,object_type)
rects_counter += 1
# loop over the tracked objects
objects, last_location, mappings = ct.update(rects)
#print(mappings, '*#*$*#$*#$*$#*')
total = []
for (objectID, centroid) in objects.items():
# draw both the ID of the object and the centroid of the
# object on the output frame
text = "ID {}".format(objectID)
cv2.putText(src, text, (centroid[0] - 10, centroid[1] - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
cv2.circle(src, (centroid[0], centroid[1]), 4, (0, 255, 0), -1)
total.append((objectID, centroid))
yield total, mappings
for (objectID, centroid_list) in last_location.items():
if len(centroid_list) == 2:
p1 = (int(centroid_list[0][0]), int(centroid_list[0][1]))
p2 = (int(centroid_list[1][0]), int(centroid_list[1][1]))
# Method # 1
total.append((p1, p2))
# Method # 2
if objectID not in total_points:
total_points[objectID] = [(p1, p2)]
color_id[objectID] = (random.randint(100, 250), random.randint(
100, 250), random.randint(150, 255))
else:
total_points[objectID].append((p1, p2))
# thickness = int(np.sqrt(32 / float(i + 1)) * 50)
# (p1,p2)
# method # 1
# for x in total:
# cv2.line(frame, x[0], x[1], (255, 0, 0), 10)
# method # 2
# print(total_points.keys(),total_points.values())
for key, value in total_points.items():
for point in value:
cv2.line(src, point[0], point[1], color_id[key], 10)
# cv2.line(frame, point[0], point[1], (0, 0, 255), 10)
# yield total_points
total = []
detections_in_frame = []
# for region in bboxes:
# values = bboxes[region]
# cv2.rectangle(src, (values[0], values[1]), (values[2], values[3]), (0, 0, 255), 2)
for line in lines:
vals = lines[line]['coords']
cv2.line(src, (vals[0], vals[1]), (vals[2], vals[3]), (255, 0, 0), 5)
#this is for station monitoring video
cv2.imshow("Frame", src)
cv2.imwrite("output/image" + str(count) + ".jpg", src)
count+=1
key = cv2.waitKey(1) & 0xFF
# if the `q` key was pressed, break from the loop
if key == ord("q"):
break
# for ret in utilspytorch():
# continue