Update app.py

app.py

@@ -5,20 +5,29 @@ import pickle
 import torch
 import time
 import pandas as pd
-import sklearn 
+import sys
+import os
+
+# Add YOLOv7 repository to the system path
+sys.path.append(os.path.join(os.getcwd(), 'yolov7'))
+
+from models.experimental import attempt_load
+from utils.general import non_max_suppression, scale_coords
+from utils.datasets import letterbox
 
 # Load ML model
 model = pickle.load(open('model.pkl', 'rb'))
 
-# Load YOLOv7 model (COCO pretrained: person=0, cellphone=67)
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-yolo_model = torch.hub.load('WongKinYiu/yolov7', 'custom', 'yolov7.pt', source='local').to(device)
+# Load YOLOv7 model
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+yolo_model = attempt_load('yolov7.pt', map_location=device)
+yolo_model.eval()
 
 # Streamlit UI setup
-st.set_page_config(page_title="YOLOv7 Attention Detector", layout='wide')
-st.title("🎥 Real-Time Attention Detector (YOLOv7 Face + Phone)")
-
+st.set_page_config(page_title="Multi-Face Attention Detector", layout='wide')
+st.title("🎥 Real-Time Multi-Face Attention Detector")
 run = st.checkbox('Start Webcam')
+
 FRAME_WINDOW = st.image([])
 attention_log = []
 start_time = time.time()
@@ -31,46 +40,28 @@ if run:
             st.warning("⚠️ Cannot access webcam.")
             break
 
-        h, w, _ = frame.shape
-
-        # YOLOv7 detection (face + phone)
-        results = yolo_model([frame])
-        faces, phones = [], []
-
-        for det in results.xyxy[0]:
-            x1, y1, x2, y2, conf, cls = map(float, det[:6])
-            cls_id = int(cls)
-            if cls_id == 0:  # person → treat as face
-                faces.append((int(x1), int(y1), int(x2 - x1), int(y2 - y1), conf))
-            if cls_id == 67:  # cellphone
-                phones.append((int(x1), int(y1), int(x2 - x1), int(y2 - y1), conf))
-
-        for i, (face_x, face_y, face_w, face_h, face_conf) in enumerate(faces):
-            pose_x, pose_y = face_x + face_w // 2, face_y + face_h // 2
-
-            phone_detected = int(any(
-                face_x < p[0] + p[2] < face_x + face_w and face_y < p[1] + p[3] < face_y + face_h
-                for p in phones
-            ))
-
-            if phones:
-                phone_x, phone_y, phone_w, phone_h, phone_conf = phones[0]
-            else:
-                phone_x, phone_y, phone_w, phone_h, phone_conf = 0, 0, 0, 0, 0
-
-            feature_vector = np.array([[1, face_x, face_y, face_w, face_h, face_conf,
-                                        0, 0, pose_x, pose_y, phone_detected,
-                                        phone_x, phone_y, phone_w, phone_h, 0.8]])
-
-            pred = model.predict(feature_vector)[0]
-            attention_text = 'Attentive' if pred == 0 else 'Inattentive'
-
-            attention_log.append({'face_id': i + 1, 'time': time.time() - start_time, 'state': attention_text})
-
-            color = (0, 255, 0) if pred == 0 else (0, 0, 255)
-            cv2.rectangle(frame, (face_x, face_y), (face_x + face_w, face_y + face_h), color, 2)
-            cv2.putText(frame, f'Face {i + 1}: {attention_text}', (face_x, face_y - 10),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
+        img = letterbox(frame, new_shape=640)[0]
+        img = img[:, :, ::-1].transpose(2, 0, 1)
+        img = np.ascontiguousarray(img)
+        img = torch.from_numpy(img).to(device)
+        img = img.float()
+        img /= 255.0
+        if img.ndimension() == 3:
+            img = img.unsqueeze(0)
+
+        # Inference
+        pred = yolo_model(img)[0]
+        pred = non_max_suppression(pred, 0.25, 0.45, classes=None, agnostic=False)
+
+        # Process detections
+        for i, det in enumerate(pred):
+            if len(det):
+                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], frame.shape).round()
+                for *xyxy, conf, cls in reversed(det):
+                    label = f'{int(cls)} {conf:.2f}'
+                    cv2.rectangle(frame, (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3])), (0, 255, 0), 2)
+                    cv2.putText(frame, label, (int(xyxy[0]), int(xyxy[1]) - 10),
+                                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
 
         FRAME_WINDOW.image(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
 
@@ -79,14 +70,14 @@ if run:
 
     cap.release()
 
-# Dashboard
+# Process log for dashboard
 if attention_log:
     df = pd.DataFrame(attention_log)
     attentive = df[df['state'] == 'Attentive'].shape[0]
     inattentive = df[df['state'] == 'Inattentive'].shape[0]
     st.markdown("### 📊 Attention Statistics")
-    st.write(f"✅ Attentive: {attentive}")
-    st.write(f"⚠️ Inattentive: {inattentive}")
+    st.write(f"✅ Attentive detections: {attentive}")
+    st.write(f"⚠️ Inattentive detections: {inattentive}")
     st.dataframe(df.tail(10))
     st.line_chart(df.groupby('time')['state'].apply(lambda x: (x == 'Attentive').mean()))
     st.download_button("Download Log as CSV", df.to_csv(index=False), file_name="attention_log.csv")