Update detection.py

detection.py

@@ -1,63 +1,73 @@
-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
+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)