Preliminary results

app.py

@@ -4,16 +4,92 @@
 
 import streamlit as st
 from datasets import load_dataset, Image
-from PIL import Image
+from torch.utils.tensorboard.summary import draw_boxes
+from transformers import pipeline
+from PIL import Image, ImageDraw, ImageFont
+import torch
 
 # Load dataset from Hugging Face
-
 dataset = load_dataset("gcesar/spinach")
 
 # Call image using datasets[vision]
+# dataset["train"][0]["image"]
+
+# Check for mps
+# torch.backends.mps.is_built()
+
+# Assign GPU
+#device = torch.device("mps")
+
+# Use GPU
+# .to(device)
+# pipeline(device=device)
+
+# Create pipeline model
+pipe = pipeline(task="object-detection", model="haiquanua/weed_detr")
+# Create pipeline model with mps
+# pipe = pipeline(task="object-detection", model="haiquanua/weed_detr", device=device)
+
+
+
+# Professor Haiquan Li function draw_boxes from haiquanua/BAT102
+def draw_boxes(im: Image.Image, preds, threshold: float = 0.25,
+               class_map={"LABEL_0": "Weed", "LABEL_1": "lettuce", "LABEL_2": "Spinach"}) -> Image.Image:
+    """Draw bounding boxes + labels on a PIL image."""
+    im = im.convert("RGB")
+    draw = ImageDraw.Draw(im)
+    try:
+        # A small default bitmap font (portable in Spaces)
+        font = ImageFont.load_default()
+    except Exception:
+        font = None
+
+    for p in preds:
+        if p.get("score", 0) < threshold:
+            continue
+        box = p["box"]  # {'xmin','ymin','xmax','ymax'}
+        class_label = class_map.get(p['label'], p['label'])
+        label = f"{class_label} {p['score']:.2f}"
+        xy = [(box["xmin"], box["ymin"]), (box["xmax"], box["ymax"])]
+
+        if p['label'] == 'LABEL_0':
+            col = (255, 0, 0)  # red
+        elif p['label'] == 'LABEL_1':
+            col = (0, 255, 0)  # green
+        else:
+            col = 'yellow'
+
+        # rectangle + label background
+        draw.rectangle(xy, outline=(255, 0, 0), width=3)
+        tw, th = draw.textlength(label, font=font), 14 if font is None else font.size + 6
+        x0, y0 = box["xmin"], max(0, box["ymin"] - th - 2)
+        draw.rectangle([x0, y0, x0 + tw + 6, y0 + th + 2], fill=(0, 0, 0))
+        draw.text((x0 + 3, y0 + 2), label, fill=(255, 255, 255), font=font)
+
+    counts = {}
+    for p in preds:
+        if p.get("score", 0) >= threshold:
+            counts[p["label"]] = counts.get(p["label"], 0) + 1
+    caption = ", ".join(f"{k}: {v}" for k, v in sorted(counts.items())) or "No detections"
+    return im
+
+
 
+# Set tittle
+st.title("Weed Detector")
 
-st.write(dataset["train"][0]["image"])
+# Iterate images
+for i in range(0, 20):
+    im = dataset["train"][i]["image"]
+    # Predict pipe
+    preds = pipe(im)
+    # Draw boxes
+    img = draw_boxes(im, preds)
+    # Display images with streamlit
+    st.write(img)
 
 
+# img = draw_boxes(im, preds)
 
+# st.write(img)
+# st.image(img)