Spaces:
Build error
Build error
| 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 | |