Spaces:
Runtime error
Runtime error
File size: 6,031 Bytes
4b9a23d b4fa442 4b9a23d 840a1e1 4b9a23d 840a1e1 b4fa442 4b9a23d 5f0de3f 4b9a23d |
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 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 |
import numpy as np
from ultralytics import YOLO
import cv2
import cvzone
import math
from sort import *
# MongoDB connection URI with a default database (replace with your actual values)
# def connectMongo() -> pymongo.database.Database:
# try:
# try:
# client = pymongo.MongoClient('mongodb+srv://INFINIX:INFINIX@cluster0.rubyoda.mongodb.net/?retryWrites=true&w=majority')
# db = client["INFINIX"]
# print(db)
# if db != None:
# print("connected to db")
# return db
# except Exception as e:
# print(e)
# except Exception as e:
# return "Error in Connecting to MongoDB" + str(e)
# db=connectMongo()
# collection=db['BUS_DETS']
# Create a collection to store entry count
# entry_count_collection = db.entry_count
cap = cv2.VideoCapture('./trialFootage.mp4')
model = YOLO("../Yolo-Weights/yolov8n.pt")
occupancy = 0
coming = 0
goin = 0
ListPeople = []
dict = {}
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
print(f"Video Resolution: {width}x{height}")
classNames = ["person", "bicycle", "car", "motorbike", "aeroplane", "bus", "train", "truck", "boat",
"traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
"dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella",
"handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat",
"baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
"fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli",
"carrot", "hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant", "bed",
"diningtable", "toilet", "tvmonitor", "laptop", "mouse", "remote", "keyboard", "cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors",
"teddy bear", "hair drier", "toothbrush"
]
# Tracking
tracker = Sort(max_age=20, min_hits=3, iou_threshold=0.3)
yelloLine = [270, 0, 270, 600]
RedLine = [173, 0, 173, 600]
totalCountUp = []
#mask=cv2.imread('mask.jpg')
while True:
success, img = cap.read()
# imgRegion=cv2.bitwise_and(img,mask)
results = model(img, stream=True)
detections = np.empty((0, 5))
for r in results:
boxes = r.boxes
for box in boxes:
# Bounding Box
x1, y1, x2, y2 = box.xyxy[0]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
# cv2.rectangle(img,(x1,y1),(x2,y2),(255,0,255),3)
w, h = x2 - x1, y2 - y1
# Confidence
conf = math.ceil((box.conf[0] * 100)) / 100
# Class Name
cls = int(box.cls[0])
currentClass = classNames[cls]
if currentClass == "person" and conf > 0.3:
# cvzone.putTextRect(img, f'{currentClass} {conf}', (max(0, x1), max(35, y1)),
# scale=0.6, thickness=1, offset=3)
# cvzone.cornerRect(img, (x1, y1, w, h), l=9, rt=5)
currentArray = np.array([x1, y1, x2, y2, conf])
detections = np.vstack((detections, currentArray))
resultsTracker = tracker.update(detections)
for result in resultsTracker:
x1, y1, x2, y2, id = result
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
print(result)
w, h = x2 - x1, y2 - y1
cvzone.cornerRect(img, (x1, y1, w, h), l=9, rt=2, colorR=(255, 0, 255))
cvzone.putTextRect(img, f' {int(id)}', (max(0, x1), max(35, y1)),
scale=2, thickness=3, offset=10)
cx, cy = x1 + w // 2, y1 + h // 2
cv2.circle(img, (cx, cy), 5, (255, 0, 255), cv2.FILLED)
if yelloLine[0] - 20 < cx < yelloLine[2] + 20:
if totalCountUp.count(id) == 0:
totalCountUp.append(id)
dict[id] = [False]
cv2.line(img, (yelloLine[0], yelloLine[1]), (yelloLine[2], yelloLine[3]), (0, 0, 255), 5)
elif totalCountUp.count(id) == 1:
if (dict[id].count(False) < 1):
dict[id].append(False)
cv2.line(img, (yelloLine[0], yelloLine[1]), (yelloLine[2], yelloLine[3]), (0, 0, 255), 5)
if RedLine[0] - 20 < cx < RedLine[2] + 30:
if totalCountUp.count(id) == 0:
totalCountUp.append(id)
dict[id] = [True]
cv2.line(img, (RedLine[0], RedLine[1]), (RedLine[2], RedLine[3]), (0, 255, 200), 5)
elif totalCountUp.count(id) == 1:
if (dict[id].count(True) < 1):
dict[id].append(True)
cv2.line(img, (RedLine[0], RedLine[1]), (RedLine[2], RedLine[3]), (0, 255, 200), 5)
print(totalCountUp)
entry_count = 0
for i in dict.values():
if (len(i) == 2):
if i[0] == True and i[1] == False:
if entry_count > 0:
entry_count -= 1
if i[0] == False and i[1] == True:
entry_count += 1
print('count is ', entry_count)
print(dict)
# # Update the MongoDB collection with the current count
# entry_count_collection.update_one({}, {"$set": {"count":entry_count}}, upsert=True)
# collection.update_one({"id": "your_document_id"}, {"$set": {"entry_count": entry_count}})
cv2.putText(img, str(entry_count), (110, 245), cv2.FONT_HERSHEY_PLAIN, 5, (50, 50, 230), 7)
print('count is ', entry_count)
print(dict)
#cv2.imshow("Image", img)
#cv2.waitKey(1)
print(entry_count)
# collection.update_one({"id":"826587"},
# {"$set": {
# "entry_count":entry_count,
# }})
|