Basmala3li/Computer_Vision_Project
0
1import io2import base643import os4from flask import Flask, request, render_template_string5from ultralytics import YOLO6from PIL import Image, ImageDraw7import torch8 9app = Flask(__name__)10 11device = "cuda" if torch.cuda.is_available() else "cpu"12 13model = YOLO("best.pt").to(device)14 15HTML_TEMPLATE = """16<!doctype html>17<html lang="en">18<head>19 <meta charset="utf-8">20 <meta name="viewport" content="width=device-width, initial-scale=1">21 <title>YOLO Object Detection</title>22 <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0/dist/css/bootstrap.min.css" rel="stylesheet">23 <style>24 body { background-color: #f5f5f5; }25 .container { margin-top: 50px; }26 .card { margin-top: 20px; }27 .btn-primary { background-color: #4CAF50; border: none; }28 </style>29</head>30<body>31<div class="container">32 <h1 class="text-center text-success">YOLO Object Detection</h1>33 <p class="text-center">Upload your image and objects will be detected</p>34 <form method="post" action="/predict" enctype="multipart/form-data" class="text-center">35 <input type="file" name="image" class="form-control mb-3" required>36 <button type="submit" class="btn btn-primary btn-lg">Detect</button>37 </form>38 39 {% if message %}40 <div class="alert alert-warning text-center mt-3">{{ message }}</div>41 {% endif %}42 43 {% if img %}44 <div class="card mx-auto" style="width: 600px;">45 <img src="data:image/png;base64,{{ img }}" class="card-img-top" alt="Result">46 </div>47 {% endif %}48</div>49</body>50</html>51"""52 53@app.route("/")54def index():55 return render_template_string(HTML_TEMPLATE)56 57@app.route("/predict", methods=["POST"])58def predict():59 file = request.files["image"]60 image = Image.open(file).convert("RGB")61 draw = ImageDraw.Draw(image)62 63 results = model(image, conf=0.25, device=device)[0]64 65 person_detected = False 66 67 for box in results.boxes:68 cls_id = int(box.cls[0])69 cls_name = model.names[cls_id]70 conf = float(box.conf[0]) * 100 71 if cls_name.lower() != "person":72 continue73 person_detected = True74 x1, y1, x2, y2 = box.xyxy[0].tolist()75 draw.rectangle([x1, y1, x2, y2], outline="red", width=3)76 draw.text((x1, y1 - 15), f"{cls_name}: {conf:.1f}%", fill="red")77 78 if not person_detected:79 message = "No persons detected in the image."80 return render_template_string(HTML_TEMPLATE, message=message)81 82 buf = io.BytesIO()83 image.save(buf, format="PNG")84 img_b64 = base64.b64encode(buf.getvalue()).decode()85 86 return render_template_string(HTML_TEMPLATE, img=img_b64)87 88if __name__ == "__main__":89 port = int(os.environ.get("PORT", 7860))90 app.run(host="0.0.0.0", port=port)91