Update detection.py

detection.py

@@ -1,73 +1,62 @@
-import gradio as gr
-from detection import ObjectDetection
-
-examples = [
-    ['test-images/plant1.jpeg', 0.23],
-    ['test-images/plant2.jpeg', 0.45],
-    ['test-images/plant3.webp', 0.43],
-]
-
-def get_predictions(img, threshold, box_color, text_color):
-    v8_results = yolov8_detector.v8_score_frame(img)
-    v8_frame = yolov8_detector.plot_bboxes(v8_results, img, float(threshold), box_color, text_color)
-    return v8_frame
-
-# Load the YOLOv8 model for plant leaf detection and classification
-yolov8_detector = ObjectDetection('Yolov8')
-
-interface = gr.Interface(
-    fn=get_predictions,
-    inputs=[
-        gr.Image(shape=(824, 824), label="Input Image"),
-        gr.Slider(maximum=1, step=0.01, value=0.4, label="Confidence Threshold", interactive=True),
-        gr.ColorPicker(label="Box Color", value="#FF8C00"),
-        gr.ColorPicker(label="Prediction Color", value="#000000"),
-    ],
-    outputs=gr.Image(label="YOLOv8 Prediction"),
-    examples=examples,
-    live=True,
-    title="Plant Leaf Detection and Classification",
-)
-
-# Custom CSS to create a dark mode appearance
-custom_css = """
-<style>
-body {
-    background-color: #222222;
-    color: #FFFFFF;
-}
-
-h1, h2, h3, h4, h5, h6 {
-    color: #FF8C00;
-}
-
-.gradio-interface {
-    border: 1px solid #FF8C00;
-    border-radius: 10px;
-    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
-}
-
-.gradio-interface > .title {
-    background-color: #FF8C00;
-    color: #FFFFFF;
-    padding: 12px;
-    border-top-left-radius: 10px;
-    border-top-right-radius: 10px;
-}
-
-.gradio-interface > .content {
-    padding: 20px;
-}
-
-.gradio-interface > .footer {
-    background-color: #FF8C00;
-    color: #FFFFFF;
-    padding: 12px;
-    border-bottom-left-radius: 10px;
-    border-bottom-right-radius: 10px;
-}
-</style>
-"""
-
-# Inject custom CSS into the interface
-interface.launch(share=False, custom_css=custom_css)
+import cv2
+import IPython
+from PIL import ImageColor
+from ultralytics import YOLO
+
+class ObjectDetection:
+    def __init__(self, model_name='Yolov8'):
+        self.model_name = model_name
+        self.model = self.load_model()
+        self.classes = self.model.names
+        self.device = 'cpu' 
+
+    def load_model(self):
+        model = YOLO(f"weights/{self.model_name}_best.pt")
+        return model
+    
+    def v8_score_frame(self, frame):
+            results = self.model(frame)
+    
+            labels = []
+            confidences = []
+            coords = []
+    
+            for result in results:
+                boxes = result.boxes.cpu().numpy()
+                
+                label = boxes.cls
+                conf = boxes.conf
+                coord = boxes.xyxy
+    
+                labels.extend(label)
+                confidences.extend(conf)
+                coords.extend(coord)
+    
+            return labels, confidences, coords
+
+    def get_coords(self, frame, row):
+        return int(row[0]), int(row[1]), int(row[2]), int(row[3])
+    
+    def class_to_label(self, x):
+        return self.classes[int(x)]
+    
+    def get_color(self, code):
+        rgb = ImageColor.getcolor(code, "RGB")
+        return rgb 
+    
+    def plot_bboxes(self, results, frame, threshold=0.5, box_color='red', text_color='white'):
+        labels, conf, coord = results
+
+        frame = frame.copy()
+        box_color = self.get_color(box_color)
+        text_color = self.get_color(text_color)
+
+        for i in range(len(labels)):
+            if conf[i] >= threshold:
+                x1, y1, x2, y2 = self.get_coords(frame, coord[i])
+                class_name = self.class_to_label(labels[i])
+
+                cv2.rectangle(frame, (x1, y1), (x2, y2), box_color, 2)
+                cv2.putText(frame, f"{class_name} - {conf[i]*100:.2f}%", (x1, y1), cv2.FONT_HERSHEY_COMPLEX, 0.5, text_color)
+        
+        return frame