Upload folder using huggingface_hub

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
-import cv2 as cv2
 import pandas as pd
-from sahi.prediction import ObjectPrediction
-from sahi.utils.cv import visualize_object_predictions, read_image
-from ultralyticsplus import YOLO 
+import os, shutil
+from PIL import Image
+from ultralyticsplus import YOLO
+
 
 # Gradio Theme
 theme = gr.themes.Soft(
@@ -16,6 +16,7 @@ theme = gr.themes.Soft(
     background_fill_primary='*neutral_100',
 )
 
+
 # Bread Prices
 bread_types = {
     "baguette": {"name": "Baguette", "price": 108},
@@ -35,41 +36,60 @@ bread_types = {
     "whole-grain-bread": {"name": "Whole Grain Bread", "price": 10},
 }
 
+
 # Instantiate the model
-model = YOLO('best.pt')
+model = YOLO("best.pt")
 
-# Bread Prediction function
-def detect_bread(image):
-    results = model.predict(image, conf=0.4)
-    result = results[0]
-    object_prediction_list = []
+
+# Converts image input into a list
+def preprocess_image(image):
+    img_list = []
+ 
+    for im in image:
+        image = Image.open(im.name)
+        img_list.append(image)
     
-    # Image Output function
-    for box in result.boxes:
-        bbox = box.xyxy[0].tolist()
-        score = round(box.conf[0].item(), 2)
-        category_id = box.cls[0]
-        category_name = result.names[box.cls[0].item()]
-
-        object_prediction = ObjectPrediction(
-            bbox=bbox,
-            category_id=int(category_id),
-            score=score,
-            category_name=category_name,
-        )
-        object_prediction_list.append(object_prediction)
+    return img_list
+
 
-    image = read_image(image)
-    output_image = visualize_object_predictions(image=image, object_prediction_list=object_prediction_list)
+# Gets all output images
+def get_predictions(directory):
+    allowed_extensions = ('.png', '.jpg', '.jpeg', '.gif', '.bmp')
+    return [
+        os.path.join(root, file)
+        for root, _, files in os.walk(directory)
+        for file in files
+        if file.lower().endswith(allowed_extensions)
+    ]
+
+
+# Clear output from previous detection
+def clear_output():
+    shutil.rmtree('output/', ignore_errors=True)
+
+
+# Bread Prediction function
+def detect_bread(image):
+    clear_output()
+    image_list = preprocess_image(image)
+    results = model.predict(image_list, conf=0.4, save=True, hide_conf=True, project = "output", name = "results")
 
-    # Receipt Output function
     detected_classes = []
+
+    for result in results:
+        for cls in result.boxes.cls:                                # Stores all detected classes in the list
+            detected_classes.append(result.names[int(cls)])   
+
+    receipt = generate_receipt(detected_classes)
+
+    return get_predictions(f'output/results'), receipt
+
+
+# Generate Receipt function
+def generate_receipt(detected_classes):
     detected_items = []
     counts = {}     # Dictionary to store bread type counts
-    
-    for cls in result.boxes.cls:                                # Stores all detected classes in the list
-        detected_classes.append(result.names[int(cls)])         
-
+ 
     for item_class in detected_classes:                         # Counts the quantity of each class
         counts[item_class] = counts.get(item_class, 0) + 1      
 
@@ -103,18 +123,27 @@ def detect_bread(image):
 
     df = pd.DataFrame(data)
 
-    return output_image['image'], df
+    return df
+
 
 # Export to CSV function
 def export_csv(df):
     df.to_csv("receipt.csv", index=False)
     return gr.File.update(value="receipt.csv", visible=True)
 
+
 # Export to JSON function
 def export_json(df):
     df.to_json("receipt.json")
     return gr.File.update(value="receipt.json", visible=True)
 
+
+# Select image from Files
+def preview(files, sd: gr.SelectData):
+    prev = files[sd.index].name
+    return gr.Image.update(value=prev, visible=True)
+
+
 # Gradio Interface
 with gr.Blocks(theme=theme) as demo:
 
@@ -124,13 +153,12 @@ with gr.Blocks(theme=theme) as demo:
 
     with gr.Row():
         with gr.Column():
-            fn = detect_bread
-            img_input = gr.Image(type="filepath", label="Input Image")
-            #img_input = gr.Files(file_types=["filepath"], label="Input Image") 
-            detect_btn = gr.Button(variant="primary", value="Detect")
-
+                fn = detect_bread
+                img_input = gr.Files(file_types=["image"], label="Input Image")
+                img_preview = gr.Image(label="Preview Image", interactive=False, visible=False)
+                detect_btn = gr.Button(variant="primary", value="Detect")
         with gr.Column():
-            img_output = gr.Image(type="filepath", label='Output Image')
+            img_output = gr.Gallery(label='Output Image')
             receipt_output = gr.Dataframe(
                     headers=["Item", "Quantity", "Price", "Amount"],
                     datatype=["str", "number", "number", "number"],
@@ -138,21 +166,22 @@ with gr.Blocks(theme=theme) as demo:
                     interactive=False,
                 )
             with gr.Row():
-                clear_btn = gr.ClearButton([img_input, img_output, receipt_output])
-                export_btn = gr.Button(variant="primary", value="Export as CSV")
+                clear_btn = gr.ClearButton()
+                export_csv_btn = gr.Button(variant="primary", value="Export as CSV")
                 export_json_btn = gr.Button(variant="primary", value="Export as JSON")
             with gr.Row():
-                csv = gr.File(interactive=False, visible=False)
-                
-    with gr.Row():
-        gr.Examples(
-            examples = ["examples/bonete.jpg", "examples/pandesal.jpg", "examples/croissant_baguette.jpg", "examples/slices.png"],
-            inputs = img_input,
-        )
+                csv = gr.File(interactive=False, visible=False)           
+ 
 
+    # Gradio Buttons
+    img_input.select(preview, img_input, img_preview)        
     detect_btn.click(detect_bread, inputs=img_input, outputs=[img_output, receipt_output])
-    export_btn.click(export_csv, receipt_output, csv)
+    export_csv_btn.click(export_csv, receipt_output, csv)
     export_json_btn.click(export_json, receipt_output, csv)
+    clear_btn.click(lambda: [None, None, None, gr.File.update(visible=False), gr.Image.update(visible=False)],
+        outputs=[img_input, img_output, receipt_output, csv, img_preview]
+        )
+
 
 demo.queue()
 demo.launch()

requirements.txt

@@ -1,5 +1,4 @@
 gradio==3.40.1
-opencv_python==4.8.0.74
 pandas==2.0.3
-sahi==0.11.14
+Pillow==10.0.0
 ultralyticsplus==0.0.28