Spaces:
Build error
Build error
File size: 4,274 Bytes
9fbe262 | 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 | import cv2
from ultralytics import YOLO
import os
from datetime import datetime
class GasBottleCounter:
def __init__(self, model_path):
self.model = YOLO(model_path)
self.IOU_THRESHOLD = 0.5
self.counts = {}
self.total = 0
def iou(self, boxA, boxB):
xA = max(boxA[0], boxB[0])
yA = max(boxA[1], boxB[1])
xB = min(boxA[2], boxB[2])
yB = min(boxA[3], boxB[3])
inter = max(0, xB - xA) * max(0, yB - yA)
if inter == 0:
return 0.0
areaA = (boxA[2]-boxA[0]) * (boxA[3]-boxA[1])
areaB = (boxB[2]-boxB[0]) * (boxB[3]-boxB[1])
return inter / float(areaA + areaB - inter)
def generate_frames(self, input_path):
"""Generator that yields processed frames and updates internal count state."""
cap = cv2.VideoCapture(input_path)
if not cap.isOpened():
return
self.counts = {}
self.total = 0
detected_ids = set()
id_to_box = {}
while True:
ret, frame = cap.read()
if not ret:
break
results = self.model.track(frame, persist=True, conf=0.4, verbose=False)
for r in results:
if r.boxes is None:
continue
for box in r.boxes:
if box.id is None:
continue
track_id = int(box.id[0])
cls_id = int(box.cls[0])
cls_name = r.names[cls_id]
conf = float(box.conf[0])
x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
curr_box = (x1, y1, x2, y2)
# Professional visualization (Blue/White)
cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 78, 3), 2) # Deep blue/orange accent
cv2.putText(
frame,
f"{cls_name} #{track_id}",
(x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 78, 3), 2
)
if track_id not in detected_ids:
is_duplicate = False
for known_id, known_box in id_to_box.items():
if known_id not in detected_ids:
continue
if self.iou(curr_box, known_box) > self.IOU_THRESHOLD:
is_duplicate = True
break
if not is_duplicate:
detected_ids.add(track_id)
self.counts[cls_name] = self.counts.get(cls_name, 0) + 1
id_to_box[track_id] = curr_box
# Build Overlay
self.total = sum(self.counts.values())
# Simple header overlay
overlay = frame.copy()
cv2.rectangle(overlay, (0, 0), (frame.shape[1], 40), (255, 255, 255), -1)
cv2.addWeighted(overlay, 0.8, frame, 0.2, 0, frame)
cv2.putText(
frame,
f"LIVE ANALYTICS | TOTAL DETECTED: {self.total}",
(20, 28),
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (3, 78, 255), 2 # Professional blue
)
# Encode frame to JPEG
ret, buffer = cv2.imencode('.jpg', frame)
frame_bytes = buffer.tobytes()
yield (b'--frame\r\n'
b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')
cap.release()
# Save results to file
os.makedirs("results", exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
report_path = os.path.join("results", f"report_{timestamp}.txt")
with open(report_path, "w") as f:
f.write(f"GazCounter Industrial Report - {timestamp}\n")
f.write("=" * 40 + "\n")
for brand, count in sorted(self.counts.items()):
f.write(f" {brand}: {count}\n")
f.write("=" * 40 + "\n")
f.write(f" TOTAL: {self.total}\n")
def get_results(self):
return self.counts, self.total
|