Spaces:
Runtime error
Runtime error
| 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, | |
| # }}) | |