Update app.py

app.py

@@ -1,5 +1,6 @@
 from fastapi import FastAPI, File, UploadFile, HTTPException
 import torch
+from ultralytics import YOLO
 import cv2
 import requests
 import json
@@ -7,96 +8,125 @@ import time
 import os
 import numpy as np
 from pathlib import Path
+from datetime import datetime
+import logging
 
 # --- Configuration ---
-RTSP_URL = "rtsp://your_rtsp_stream_url"
-SALESFORCE_URL = "https://your_salesforce_instance_url"
-SALESFORCE_TOKEN = "your_salesforce_access_token"
-HUGGINGFACE_API_URL = "https://huggingface.co/your_model_endpoint"
-HUGGINGFACE_TOKEN = "your_huggingface_api_token"
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+RTSP_URL = os.getenv("RTSP_URL", "rtsp://localhost:8554/stream")
+SALESFORCE_URL = os.getenv("SALESFORCE_URL", "https://your_salesforce_instance_url")
+SALESFORCE_TOKEN = os.getenv("SALESFORCE_TOKEN", "your_salesforce_access_token")
+HUGGINGFACE_API_URL = os.getenv("HUGGINGFACE_API_URL", "https://api-inference.huggingface.co/models/PrashanthB461/SafetyViolationAI1")
+HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN", "your_huggingface_api_token")
 
 # --- Initialize FastAPI app ---
 app = FastAPI()
 
 # --- YOLOv8 Model ---
 class YOLOv8Model:
-    def __init__(self, model_name='yolov8'):
-        self.model = torch.hub.load('ultralytics/yolov5', model_name)  # YOLOv8 based on YOLOv5
-        self.model.eval()
+    def __init__(self, model_path='yolov8n.pt'):
+        try:
+            self.model = YOLO(model_path)  # Load YOLOv8 model
+            logger.info("YOLOv8 model loaded successfully")
+        except Exception as e:
+            logger.error(f"Failed to load YOLOv8 model: {e}")
+            raise
 
     def predict(self, image):
-        results = self.model(image)  # Inference
-        return results.pandas().xywh[0]  # Bounding boxes, class names, confidence score
-
+        try:
+            results = self.model(image)  # Inference
+            return results.pandas().xyxy[0]  # Bounding boxes, class names, confidence score
+        except Exception as e:
+            logger.error(f"Prediction error: {e}")
+            raise
 
 # --- Preprocessing RTSP Frame ---
 def preprocess_frame(frame):
-    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # Convert frame to RGB
-    img_resized = cv2.resize(img, (640, 640))  # Resize to 640x640 (YOLOv8 input size)
-    return img_resized
-
+    try:
+        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        img_resized = cv2.resize(img, (640, 640))
+        return img_resized
+    except Exception as e:
+        logger.error(f"Frame preprocessing error: {e}")
+        raise
 
 # --- RTSP Stream Handler ---
 def capture_rtsp_frames(rtsp_url: str):
-    cap = cv2.VideoCapture(rtsp_url)
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if ret:
-            timestamp = time.time()
-            yield frame, timestamp
-        else:
-            break
-    cap.release()
-
+    try:
+        cap = cv2.VideoCapture(rtsp_url)
+        if not cap.isOpened():
+            logger.error(f"Failed to open RTSP stream: {rtsp_url}")
+            raise ValueError("RTSP stream not accessible")
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if ret:
+                timestamp = datetime.utcnow().isoformat()
+                yield frame, timestamp
+            else:
+                logger.warning("Failed to read frame from RTSP stream")
+                break
+        cap.release()
+    except Exception as e:
+        logger.error(f"RTSP capture error: {e}")
+        raise
 
 # --- Save Violations and Snapshots ---
 def save_snapshot(frame):
-    filename = f"snapshot_{int(time.time())}.jpg"
-    snapshot_path = Path("./snapshots") / filename
-    os.makedirs("./snapshots", exist_ok=True)
-    cv2.imwrite(str(snapshot_path), frame)
-    return f"http://localhost/snapshots/{filename}"  # URL for local testing
-
+    try:
+        filename = f"snapshot_{int(time.time())}.jpg"
+        snapshot_path = Path("/snapshots") / filename
+        os.makedirs("/snapshots", exist_ok=True)
+        cv2.imwrite(str(snapshot_path), frame)
+        return f"/snapshots/{filename}"  # Relative path for containerized env
+    except Exception as e:
+        logger.error(f"Snapshot saving error: {e}")
+        raise
 
 def log_violation(violation_data):
-    log_file = Path("./violation_logs.json")
-    if log_file.exists():
-        with open(log_file, "r") as f:
-            logs = json.load(f)
-    else:
+    try:
+        log_file = Path("/snapshots/violation_logs.json")
         logs = []
-
-    logs.append(violation_data)
-
-    with open(log_file, "w") as f:
-        json.dump(logs, f, indent=4)
-
+        if log_file.exists():
+            with open(log_file, "r") as f:
+                logs = json.load(f)
+        logs.append(violation_data)
+        with open(log_file, "w") as f:
+            json.dump(logs, f, indent=4)
+    except Exception as e:
+        logger.error(f"Violation logging error: {e}")
+        raise
 
 # --- Notification System ---
 def send_alert(violation):
-    print(f"Alert! {violation['violation_type']} detected. Severity: {violation['severity']}")
-    # Implement your notification logic (email/SMS) here
-
+    logger.info(f"Alert! {violation['violation_type']} detected. Severity: {violation['severity']}")
+    # Placeholder for email/SMS notification logic
 
 # --- Salesforce Integration ---
 def create_salesforce_violation_record(violation_data):
-    salesforce_url = f"{SALESFORCE_URL}/services/data/vXX.0/sobjects/Safety_Violation_Log__c/"
-    headers = {
-        'Authorization': f'Bearer {SALESFORCE_TOKEN}',
-        'Content-Type': 'application/json'
-    }
-    violation_obj = {
-        'Site_ID__c': violation_data['site_id'],
-        'Violation_Type__c': violation_data['violation_type'],
-        'Timestamp__c': violation_data['timestamp'],
-        'Snapshot_URL__c': violation_data['snapshot_url'],
-        'Severity__c': violation_data['severity'],
-        'Alert_Sent__c': True,
-        'Resolved__c': False
-    }
-    response = requests.post(salesforce_url, headers=headers, data=json.dumps(violation_obj))
-    return response.json()
-
+    try:
+        salesforce_url = f"{SALESFORCE_URL}/services/data/v60.0/sobjects/Safety_Violation_Log__c/"
+        headers = {
+            'Authorization': f'Bearer {SALESFORCE_TOKEN}',
+            'Content-Type': 'application/json'
+        }
+        violation_obj = {
+            'Site_ID__c': violation_data['site_id'],
+            'Camera_ID__c': violation_data['camera_id'],
+            'Violation_Type__c': violation_data['violation_type'],
+            'Timestamp__c': violation_data['timestamp'],
+            'Snapshot_URL__c': violation_data['snapshot_url'],
+            'Severity__c': violation_data['severity'],
+            'Alert_Sent__c': True,
+            'Resolved__c': False
+        }
+        response = requests.post(salesforce_url, headers=headers, data=json.dumps(violation_obj))
+        response.raise_for_status()
+        return response.json()
+    except Exception as e:
+        logger.error(f"Salesforce integration error: {e}")
+        raise
 
 # --- API Routes ---
 @app.post("/detect_violation/")
@@ -108,21 +138,24 @@ async def detect_violation():
             results = model.predict(frame_processed)
 
             for index, row in results.iterrows():
+                severity = "Critical" if row['conf'] > 0.8 else "Moderate" if row['conf'] > 0.5 else "Minor"
                 violation = {
-                    'site_id': "Site1",  # Placeholder, should be dynamic
+                    'site_id': "Site1",
+                    'camera_id': "Camera1",
                     'violation_type': row['name'],
                     'timestamp': timestamp,
                     'snapshot_url': save_snapshot(frame),
-                    'severity': row['confidence']
+                    'severity': severity
                 }
-                create_salesforce_violation_record(violation)  # Log to Salesforce
-                send_alert(violation)  # Send alert to site HSE
+                log_violation(violation)
+                create_salesforce_violation_record(violation)
+                send_alert(violation)
 
         return {"status": "Violation detection complete."}
     except Exception as e:
+        logger.error(f"Error processing stream: {e}")
         raise HTTPException(status_code=500, detail=f"Error processing stream: {e}")
 
-
 @app.post("/upload_image/")
 async def upload_image(file: UploadFile = File(...)):
     try:
@@ -132,9 +165,20 @@ async def upload_image(file: UploadFile = File(...)):
         results = model.predict(image)
         return {"results": results.to_dict()}
     except Exception as e:
+        logger.error(f"Error processing image: {e}")
         raise HTTPException(status_code=500, detail=f"Error processing image: {e}")
 
-
 @app.get("/health_check/")
 async def health_check():
-    return {"status": "Running smoothly"}
+    return {"status": "Running smoothly"}
+
+@app.on_event("startup")
+async def startup_event():
+    logger.info("FastAPI application starting up")
+    # Initialize any resources (e.g., check RTSP connection, model load)
+    try:
+        model = YOLOv8Model()
+        logger.info("Startup: YOLOv8 model initialized")
+    except Exception as e:
+        logger.error(f"Startup error: {e}")
+        raise