koti-malla/object_detection

1import os2from flask import Flask, render_template, request, redirect, url_for,send_from_directory3import cv24import numpy as np5from transformers import DetrImageProcessor, DetrForObjectDetection6from torchvision.transforms import functional as F7from ultralytics import YOLO8import torch9 10 11 12app = Flask(__name__)13UPLOAD_FOLDER = 'uploads'14ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}15 16app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER17 18def allowed_file(filename):19    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS20 21 22 23@app.route('/uploads/<filename>')24def uploaded_file(filename):25    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)26 27 28@app.route('/', methods=['GET', 'POST'])29def index():30    annotated_image_url = None31 32    if request.method == 'POST':33 34        # Load the YOLOv8 model35        yolo_model = YOLO('yolo/best.pt')36 37        # Load the DETR model38        processor = DetrImageProcessor.from_pretrained("detr")39        model = DetrForObjectDetection.from_pretrained("detr")40 41        # Check if a file is selected42        if 'image' not in request.files:43            return redirect(request.url)44        45        image = request.files['image']46        47        # Check if the file has a valid extension48        if image and allowed_file(image.filename):49            constant_filename = 'my_uploaded_image.jpg'  # Specify the constant name50            filename = os.path.join(app.config['UPLOAD_FOLDER'], constant_filename)51            image.save(filename)52 53            # Load the image for processing54            image = cv2.imread(filename)55 56            # Perform YOLO object detection and annotation57            yolo_results = yolo_model(image, save=False)58            yolo_image = image.copy()59            yolo_names=yolo_results[0].names60            for row in yolo_results[0].boxes.data:61                x1, y1, x2, y2, score, class_id = row.tolist()62                x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])63 64                class_name = yolo_names.get(int(class_id), 'Unknown')65                label_text = f"Class: {class_name}, Score: {score:.2f}"66                box_color = (0, 0, 255)67                label_color = (255, 255, 255)68 69                cv2.rectangle(yolo_image, (x1, y1), (x2, y2), box_color, thickness=2)70                label_size = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)[0]71                label_bottom_left = (x1, y1 - 5)72                label_top_right = (label_bottom_left[0] + label_size[0], label_bottom_left[1] - label_size[1])73                cv2.rectangle(yolo_image, label_bottom_left, label_top_right, box_color, cv2.FILLED)74                cv2.putText(yolo_image, label_text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, label_color, 1, cv2.LINE_AA)75 76 77 78 79            annotated_filename = 'annotated_my_uploaded_image.jpg'80            annotated_filepath = os.path.join(app.config['UPLOAD_FOLDER'], annotated_filename)81            cv2.imwrite(annotated_filepath, yolo_image)82            annotated_image_url = url_for('uploaded_file', filename=annotated_filename)83 84 85 86 87            88# Process the image using the processor89            inputs = processor(images=image, return_tensors="pt")90            outputs = model(**inputs)91 92            # Convert outputs (bounding boxes and class logits) to COCO API format93            # Let's only keep detections with score > 0.994            target_sizes = torch.tensor([image.shape[:2:]])95            results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.3)[0]96 97            # Convert PIL image to NumPy array for OpenCV98            #image_np = np.array(image)99            #image_cv2 = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)100            image_cv2 = image.copy()101 102            # Define the font for labels103            font = cv2.FONT_HERSHEY_SIMPLEX104            font_scale = 0.5105            font_thickness = 1106            font_color = (255, 255, 255)  # White color107 108            # Iterate over the results and draw bounding boxes and labels using OpenCV109            for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):110                box = [round(i, 2) for i in box.tolist()]111 112                # Draw the bounding box113                box = [int(b) for b in box]  # Convert to integers for drawing114                cv2.rectangle(image_cv2, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 2)  # Red rectangle115 116                # Draw the label117                label_text = f"{model.config.id2label[label.item()]}: {round(score.item(), 3)}"118                label_size = cv2.getTextSize(label_text, font, font_scale, font_thickness)[0]119                label_bottom_left = (box[0], box[1] - 5)  # Adjust label position120                label_top_right = (label_bottom_left[0] + label_size[0], label_bottom_left[1] - label_size[1])121                cv2.rectangle(image_cv2, label_bottom_left, label_top_right, (0, 0, 255), cv2.FILLED)  # Red filled rectangle122                cv2.putText(image_cv2, label_text, (box[0], box[1] - 5), font, font_scale, font_color, font_thickness, cv2.LINE_AA)123 124 125            annotated_filename = 'dert_annotated_my_uploaded_image.jpg'126            annotated_filepath = os.path.join(app.config['UPLOAD_FOLDER'], annotated_filename)127            cv2.imwrite(annotated_filepath, image_cv2)128            dertannotated_image_url = url_for('uploaded_file', filename=annotated_filename)129 130 131 132 133 134            return render_template('index.html', image1=annotated_image_url ,image2= dertannotated_image_url)135 136 137 138 139 140 141 142 143 144        145    return render_template('index.html', image1=annotated_image_url,image2=annotated_image_url)146 147 148 149if __name__ == '__main__':150    app.run(debug=True,port=7860)