Update http_server_ir_task2.py

http_server_ir_task2.py

@@ -1,91 +1,220 @@
-from flask import Flask, request, jsonify
+import cv2
 import os
+import time
 import importlib.util
-from datetime import datetime
-from predict_task2 import Predictor
-from id_mapping import mapping
-#from show_annotation import start_annotation_process
-#from multiprocessing import Process, Queue
-#import cv2
-from show_stitched import *
-
-app = Flask(__name__)
+import supervision as sv
+from ultralytics import YOLO
 
 config_dir = os.path.abspath(os.path.dirname(__file__))
 config_path = os.path.join(config_dir, 'PC_CONFIG.py')
 spec = importlib.util.spec_from_file_location("PC_CONFIG", config_path)
 PC_CONFIG = importlib.util.module_from_spec(spec)
 spec.loader.exec_module(PC_CONFIG)
+dir = str(os.path.join(PC_CONFIG.BASE_DIR, "weights", "best_task2.pt"))
+
+class Predictor:
+    def __init__(self):
+        # Load a pre-trained yolov8n model
+        print("dir:",dir)
+        self.model = YOLO(dir) # replace model here
+    #     self.print_class_ids()  # Print class IDs upon initialization
+
+    # def print_class_ids(self):
+    #     # Print all class names and their corresponding IDs
+    #     for id, name in enumerate(self.model.names):
+    #         print(f"ID: {id}, Name: {name}")
+
+    # def predict_id(self, image_file_path, task_type):
+    #     # Load the image
+    #     image = cv2.imread(image_file_path)
+
+    #     # Run inference on the image
+    #     results = self.model(image)
+
+    #     # Print results
+    #     print(results)
+    #     # Show annotation
+    #     self.show_annotation(image, results)
+
+    #     # Extract class name
+    #     class_name, largest_size, detection_id = None, -1, None
+    #     for result in results:  # Assuming 'results' is a list
+    #         print(f"task_type is {task_type}")
+
+    #         if task_type == "TASK_2":
+    #             for prediction in result.predictions:
+    #                 print(prediction)
+    #                 class_name = prediction.class_name
+    #                 detection_id = prediction.detection_id
+    #                 if class_name != "Bullseye":
+    #                     break
+    #         else:
+    #             for prediction in result.predictions:
+    #                 print(prediction)
+    #                 if largest_size == -1 or max(prediction.width, prediction.height) > largest_size:
+    #                     largest_size = max(prediction.width, prediction.height)
+    #                     class_name = prediction.class_name
+    #                     detection_id = prediction.detection_id
+
+    #     if class_name:
+    #         print("class_name = " + class_name)
+    #     else:
+    #         print("class_name = None")
+
+    #     return class_name, results, detection_id
+
+    def predict_id(self, image_file_path, task_type):
+        # Load the image
+        image = cv2.imread(image_file_path)
+        # Validation for image existence
+        if image is None:
+            print(f"Error: Could not read image at {image_file_path}")
+            return None, None, None
+
+        # Check the image size and resize if necessary
+        if image.shape[0] != 640 or image.shape[1] != 640:
+            image = cv2.resize(image, (640, 640))  # Resize to 640x640
+
+        # Run inference on the image
+        results = self.model(image)  # Directly pass the image
+
+        # Print results
+        print(results)
+
+        # Show annotation (using YOLOv8's plotting capabilities)
+        # results[0].show()
+
+        # Extract class name, largest size, and detection ID
+        class_name, largest_size, detection_id = None, -1, 0
+        
+        # Check if there are any detections
+        if results[0].boxes is None or len(results[0].boxes) == 0:
+            print("No detections found in the image")
+            return class_name, results, detection_id
+            
+        boxes = results[0].boxes.xyxy  # Get bounding boxes (x1, y1, x2, y2)
+        scores = results[0].boxes.conf  # Get confidence scores
+        class_ids = results[0].boxes.cls  # Get class IDs
+
+        # Store all detections with their priority
+        detections_list = []
+        
+        for i in range(len(boxes)):
+            detected_class = results[0].names[int(class_ids[i])]
+            confidence = float(scores[i])
+            yolo_class_id = int(class_ids[i])
+            print(f"Processing detection {i}: {detected_class} (confidence: {confidence:.2f}, class_id: {yolo_class_id})")
+
+            if task_type == "TASK_2":
+                # Check by class name - only set Bullseye to lowest priority
+                if detected_class.lower() != 'bullseye':
+                    # All non-bullseye detections have equal priority (0), sorted by confidence
+                    print(f"  Added to list: {detected_class} with normal priority")
+                    detections_list.append({
+                        'index': i,
+                        'class_name': detected_class,
+                        'confidence': confidence,
+                        'priority': 0  # Equal priority for all non-bullseye detections
+                    })
+                else:
+                    print(f"  Bullseye detected - adding with lowest priority")
+                    detections_list.append({
+                        'index': i,
+                        'class_name': detected_class,
+                        'confidence': confidence,
+                        'priority': -10  # Lowest priority for bullseye
+                    })
+            else:
+                # Determine the largest bounding box
+                box_width = boxes[i][2] - boxes[i][0]
+                box_height = boxes[i][3] - boxes[i][1]
+                size = max(box_width, box_height)
+                detection_id = i
+
+                if largest_size == -1 or size > largest_size:
+                    largest_size = size
+                    class_name = detected_class
+                    detection_id = i
+        
+        # For TASK_2, select detection based on priority, then confidence
+        if task_type == "TASK_2" and detections_list:
+            print(f"\nTotal detections found: {len(detections_list)}")
+            # Sort by priority (descending), then by confidence (descending)
+            detections_list.sort(key=lambda x: (x['priority'], x['confidence']), reverse=True)
+            
+            # Print sorted list for debugging
+            print("Sorted detections:")
+            for det in detections_list:
+                print(f"  - {det['class_name']}: priority={det['priority']}, confidence={det['confidence']:.2f}")
+            
+            # Select the highest priority detection
+            selected = detections_list[0]
+            class_name = selected['class_name']
+            detection_id = selected['index']
+            print(f"\n✓ Selected detection: {class_name} (priority: {selected['priority']}, confidence: {selected['confidence']:.2f})")
+
+        if class_name:
+            print("class_name = " + class_name)
+            timestamp = int(time.time())
+            # Save the annotated image
+            try:
+                results[detection_id].save(f'../data/annotated_images/{class_name}_{timestamp}.jpg')
+            except:
+                print("error in saving photo!")
+        else:
+            print("class_name = None")
+
+        return class_name, results, detection_id
+
+
+    # def show_annotation(self, image, results):
+    #     # Create supervision annotators
+    #     bounding_box_annotator = sv.BoundingBoxAnnotator()
+    #     label_annotator = sv.LabelAnnotator()
+
+    #     # Process results from YOLOv8
+    #     detections = []
+    #     for result in results:
+    #         for detection in result.boxes.data:  # Accessing YOLOv8's box data
+    #             class_id = int(detection[5])  # Class ID
+    #             x1, y1, x2, y2 = map(int, detection[:4])  # Bounding box coordinates
+    #             score = float(detection[4])  # Confidence score
+
+    #             # Add to detections
+    #             detections.append({
+    #                 "bbox": [x1, y1, x2, y2],
+    #                 "confidence": score,
+    #                 "class_id": class_id
+    #             })
+
+    #     # Convert detections to the expected format for supervision
+    #     if detections:
+    #         detections = sv.Detections(
+    #             xyxy=[d["bbox"] for d in detections],
+    #             confidence=[d["confidence"] for d in detections],
+    #             class_id=[d["class_id"] for d in detections]
+    #         )
+
+    #         # Annotate the image with inference results
+    #         annotated_image = bounding_box_annotator.annotate(scene=image, detections=detections)
+    #         annotated_image = label_annotator.annotate(scene=annotated_image, detections=detections)
+
+    #         # Display the annotated image
+    #         try:
+    #             cv2.imshow("Annotated Image", annotated_image)
+    #             cv2.waitKey(0)  # Wait indefinitely until a key is pressed
+    #         except Exception as e:
+    #             print(f"Error displaying image: {e}")
+    #         finally:
+    #             cv2.destroyAllWindows()  # Close all OpenCV windows
+    #     else:
+    #         print("No detections found.")
 
-HOST = PC_CONFIG.HOST
-PORT = PC_CONFIG.IMAGE_REC_PORT
-UPLOAD_FOLDER =  os.path.join(PC_CONFIG.FILE_DIRECTORY,"image-rec","images")
-app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 
-def process_file(file_path, direction, task_type):
+if __name__ == "__main__":
+    # Example usage
     predictor = Predictor()
-    print("File received and saved successfully.")
-    print(f"Direction received: {direction}")
-    print(f"Task type received: {task_type}")
-
-    startTime = datetime.now()
-    class_name, results, detection_id = predictor.predict_id(file_path, task_type)  # Perform prediction
-    #show_annotation_queue.put((file_path, results, detection_id))
-    class_id = str(mapping.get(class_name, -1))
-    endTime = datetime.now()
-    totalTime = (endTime - startTime).total_seconds()
-    print(f"Predicted ID: {class_id}")
-    print(f"Time taken for Predicting Image = {totalTime} s")
-    return class_id
-
-@app.route('/status', methods=['GET'])
-def server_status():
-    return jsonify({'status': 'OK'})
-
-@app.route('/upload', methods=['POST'])
-def upload_file():
-    if 'file' not in request.files:
-        return jsonify({'error': 'No file part'}), 400
-    file = request.files['file']
-    direction = request.form['direction']
-    task_type = request.form['task_type']
-    if file.filename == '':
-        return jsonify({'error': 'No selected file'}), 400
-    if file:
-        filename = os.path.basename(file.filename)
-        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
-        file.save(file_path)
-        # image = cv2.imread(file_path)
-        # cv2.imshow("Uploaded Image", image)
-        # cv2.waitKey(0)  # Wait until a key is pressed
-        # cv2.destroyAllWindows()  # Close the window
-
-        # Process the file and predict
-        class_id = process_file(file_path, direction, task_type)
-        return jsonify({'message': 'File successfully uploaded', 'predicted_id': class_id}), 200
-
-@app.route('/display_stitched', methods=['POST'])
-def display_stitched():
-    showAnnotatedStitched()
-    return jsonify({'display_stitched': 'OK'})
-
-if __name__ == '__main__':
-    # show_annotation_queue = Queue()
-    # process = Process(target=start_annotation_process, args=(show_annotation_queue,))
-    # process.start()
-
-    print()
-    print(f"UPLOAD FOLDER: {UPLOAD_FOLDER}")
-    # Port 5000 if free
-    '''
-    try:
-        app.run(host=HOST, port=PORT, debug=False)
-    except:
-        print('Unable to Connect to PC_CONFIG Host and Port. Switching to 0.0.0.0:4000.')
-        app.run(host='0.0.0.0', port=4000, debug=True)
-    '''
-
-    # Run on Port 4000
-    app.run(host='0.0.0.0', port=4000, debug=True)
-
-    #process.join()
+    # Specify the path to your image
+    image_file_path = os.path.join(PC_CONFIG.FILE_DIRECTORY, "image-rec", "sample_images", "IMG_9325.jpg")
+    # Predict and display the class name
+    predictor.predict_id(image_file_path, "TASK_1")