initial commit

backend/app.py

@@ -0,0 +1,77 @@
+from flask import Flask, request, jsonify, send_file
+from pathlib import Path
+from ultralytics import YOLO
+import cv2
+import numpy as np
+import os
+import uuid
+import logging
+from io import BytesIO
+
+app = Flask(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# Initialize detector
+MODEL_PATH = str(Path(__file__).parent.parent / "runs" / "detect" / "train" / "weights" / "best.pt")
+model = YOLO(MODEL_PATH)
+
+
+@app.route('/')
+def index():
+    return send_file('test.html')
+
+
+@app.route('/detect', methods=['POST'])
+def detect():
+    if 'image' not in request.files:
+        return jsonify({'error': 'No image provided'}), 400
+
+    file = request.files['image']
+    if file.filename == '':
+        return jsonify({'error': 'No selected file'}), 400
+
+    try:
+        # Read image directly from memory
+        img_bytes = file.read()
+        img = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), cv2.IMREAD_COLOR)
+
+        # Run prediction
+        results = model.predict(img, imgsz=416, conf=0.3)
+        
+        # Generate annotated image
+        annotated = results[0].plot(line_width=2, font_size=0.5)
+        
+        # Save to results folder
+        output_dir = Path("static/results")
+        output_dir.mkdir(exist_ok=True)
+        filename = f"{uuid.uuid4()}.jpg"
+        output_path = output_dir / filename
+        cv2.imwrite(str(output_path), annotated)
+        
+        # Extract detections
+        detections = []
+        for box in results[0].boxes:
+            detections.append({
+                'box': box.xyxy[0].tolist(),
+                'confidence': float(box.conf[0]),
+                'class': int(box.cls[0])
+            })
+        
+        return jsonify({
+            'detections': detections,
+            'result_image': f"/results/{filename}"
+        })
+        
+    except Exception as e:
+        logging.error(f"Detection error: {str(e)}")
+        return jsonify({'error': str(e)}), 500
+
+@app.route('/results/<filename>')
+def get_result(filename):
+    return send_file(Path("static/results") / filename)
+
+if __name__ == '__main__':
+    # Create directories
+    Path("static/uploads").mkdir(parents=True, exist_ok=True)
+    Path("static/results").mkdir(parents=True, exist_ok=True)
+    app.run(host='0.0.0.0', port=5000)