initial commit

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,6 +1,6 @@
 ---
 title: Image Recognition-CZ3004 SC2079 Multidisciplinary Project-NTU SG
-emoji: 🐢
+emoji: 🔣
 colorFrom: indigo
 colorTo: green
 sdk: gradio

Weights/Week_8.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1cdc1c277dbbc29901e2463616e19fd74592eefca3224be23d4f54ffac2c69e9
+size 57300618

Weights/Week_9.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e6a050ce6e9fe9711d1a41aed0934ece185611d7081b12eb8eab0332c7a71307
+size 14318653

app.py

@@ -0,0 +1,239 @@
+import gradio as gr
+import torch
+from PIL import Image
+import json
+import numpy as np
+import cv2
+week8_model = torch.hub.load(
+    './', 'custom', path='Weights/Week_8.pt', source='local')
+week9_model = torch.hub.load(
+    './', 'custom', path='Weights/Week_9.pt', source='local')
+
+def draw_own_bbox(img,x1,y1,x2,y2,label,color=(36,255,12),text_color=(0,0,0)):
+    """
+    Draw bounding box on the image with text label and save both the raw and annotated image in the 'own_results' folder
+
+    Inputs
+    ------
+    img: numpy.ndarray - image on which the bounding box is to be drawn
+
+    x1: int - x coordinate of the top left corner of the bounding box
+
+    y1: int - y coordinate of the top left corner of the bounding box
+
+    x2: int - x coordinate of the bottom right corner of the bounding box
+
+    y2: int - y coordinate of the bottom right corner of the bounding box
+
+    label: str - label to be written on the bounding box
+
+    color: tuple - color of the bounding box
+
+    text_color: tuple - color of the text label
+
+    Returns
+    -------
+    None
+
+    """
+    name_to_id = {
+        "NA": 'NA',
+        "Bullseye": 10,
+        "One": 11,
+        "Two": 12,
+        "Three": 13,
+        "Four": 14,
+        "Five": 15,
+        "Six": 16,
+        "Seven": 17,
+        "Eight": 18,
+        "Nine": 19,
+        "A": 20,
+        "B": 21,
+        "C": 22,
+        "D": 23,
+        "E": 24,
+        "F": 25,
+        "G": 26,
+        "H": 27,
+        "S": 28,
+        "T": 29,
+        "U": 30,
+        "V": 31,
+        "W": 32,
+        "X": 33,
+        "Y": 34,
+        "Z": 35,
+        "Up": 36,
+        "Down": 37,
+        "Right": 38,
+        "Left": 39,
+        "Up Arrow": 36,
+        "Down Arrow": 37,
+        "Right Arrow": 38,
+        "Left Arrow": 39,
+        "Stop": 40
+    }
+    # Reformat the label to {label name}-{label id}
+    label = label + "-" + str(name_to_id[label])
+    # Convert the coordinates to int
+    x1 = int(x1)
+    x2 = int(x2)
+    y1 = int(y1)
+    y2 = int(y2)
+
+    # Save the raw image
+    # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+    # Draw the bounding box
+    img = cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
+    # For the text background, find space required by the text so that we can put a background with that amount of width.
+    (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)
+    # Print the text  
+    img = cv2.rectangle(img, (x1, y1 - 20), (x1 + w, y1), color, -1)
+    img = cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, text_color, 1)
+    return img
+
+def yolo(img, model, toggles, signal, size=1024):
+    if model == "Week 8":
+        model = week8_model
+    else:
+        model = week9_model
+
+    results = model(img)
+
+    # Original output image and results
+    original_results = json.loads(results.pandas().xyxy[0].to_json(orient="records"))
+    output_image = Image.fromarray(results.render()[0])
+
+    # Convert the results to a pandas dataframe and calculate the height and width of the bounding box and the area of the bounding box
+    df_results = results.pandas().xyxy[0]
+    df_results['bboxHt'] = df_results['ymax'] - df_results['ymin']
+    df_results['bboxWt'] = df_results['xmax'] - df_results['xmin']
+    df_results['bboxArea'] = df_results['bboxHt'] * df_results['bboxWt']
+
+    # Label with largest bbox height will be last
+    df_results = df_results.sort_values('bboxArea', ascending=False)
+
+    # Filter out Bullseye
+    pred_list = df_results 
+    if 'Ignore Bullseye' in toggles:
+        pred_list = pred_list[pred_list['name'] != 'Bullseye']
+
+    if len(pred_list) == 0:
+        return [output_image, original_results, output_image, original_results]
+
+    elif len(pred_list) == 1:
+        pred = pred_list.iloc[0]
+
+    else:
+        pred_shortlist = []
+        current_area = pred_list.iloc[0]['bboxArea']
+
+        # For each prediction, check if the confidence is greater than 0.5 and if the area is greater than 80% of the current area or 60% if the prediction is 'One'
+        for _, row in pred_list.iterrows():
+            if row['confidence'] > 0.5 and ((current_area * 0.8 <= row['bboxArea']) or (row['name'] == 'One' and current_area * 0.6 <= row['bboxArea'])):
+                # Add the prediction to the shortlist
+                pred_shortlist.append(row)
+                # Update the current area to the area of the prediction
+                current_area = row['bboxArea']
+        
+        # If only 1 prediction remains after filtering by confidence and area
+        if len(pred_shortlist) == 1:
+            # Choose that prediction
+            pred = pred_shortlist[0]
+
+        # If multiple predictions remain after filtering by confidence and area
+        else:
+            # Use signal of {signal} to filter further 
+            
+            # Sort the predictions by xmin
+            pred_shortlist.sort(key=lambda x: x['xmin'])
+
+            # If signal is 'L', choose the first prediction in the list, i.e. leftmost in the image
+            if signal == 'L':
+                pred = pred_shortlist[0]
+            
+            # If signal is 'R', choose the last prediction in the list, i.e. rightmost in the image
+            elif signal == 'R':
+                pred = pred_shortlist[-1]
+            
+            # If signal is 'C', choose the prediction that is central in the image
+            else:
+                # Loop through the predictions shortlist
+                for i in range(len(pred_shortlist)):
+                    # If the xmin of the prediction is between 250 and 774, i.e. the center of the image, choose that prediction
+                    if pred_shortlist[i]['xmin'] > 250 and pred_shortlist[i]['xmin'] < 774:
+                        pred = pred_shortlist[i]
+                        break
+                
+                # If no prediction is central, choose the one with the largest area
+                if isinstance(pred,str):
+                    # Choosing one with largest area if none are central
+                    pred_shortlist.sort(key=lambda x: x['bboxArea']) 
+                    pred = pred_shortlist[-1]
+    
+    
+    filtered_img = draw_own_bbox(np.array(img), pred['xmin'], pred['ymin'], pred['xmax'], pred['ymax'], pred['name'])
+    return [output_image, original_results, filtered_img, json.loads(pred.to_json(orient="records"))]
+
+
+inputs = [gr.inputs.Image(type='pil', label="Original Image"), 
+          gr.inputs.Radio(['Week 8', 'Week 9'], type="value", default='Week 8', label='Model Selection'),
+          gr.CheckboxGroup(["Ignore Bullseye", "Biggest BBox Only and Position-Based Heuristics",], value=["Ignore Bullseye", "Biggest BBox Only and Position-Based Heuristics"], label="Heuristic Toggles"),
+          gr.inputs.Radio(['Left', 'Center', 'Right', 'Disabled'], type="value", default='Center', label='Position Heuristic'),
+          ]
+outputs = [gr.outputs.Image(type="pil", label="Output Image"),
+           gr.outputs.JSON(label="Output JSON"),
+           gr.outputs.Image(type="pil", label="Filtered Output Image"),
+           gr.outputs.JSON(label="Filtered Output JSON")
+           ]
+
+examples = [['Examples/One.jpg'], ['Examples/Two.jpg'], ['Examples/Three.jpg'], ['Examples/1.jpg'], ['Examples/2.jpg'], ['Examples/3.jpg'], ['Examples/4.jpg'], ['Examples/5.jpg'], ['Examples/6.jpg'],
+            ['Examples/7.jpg'], ['Examples/8.jpg'], ['Examples/9.jpg'], ['Examples/10.jpg'], ['Examples/11.jpg'], ['Examples/12.jpg']]
+
+
+with gr.Blocks(css="#custom_header {min-height: 2rem; text-align: center} #custom_title {min-height: 2rem}") as demo:
+    gr.Markdown("# YOLOv5 Symbol Recognition for CZ3004/SC2079 Multi-Disciplinary Project", elem_id="custom_header")
+    gr.Markdown("Gradio Demo for YOLOv5 Symbol Recognition for CZ3004 Multi-Disciplinary Project. To use it, simply upload your image, or click one of the examples to load them.", elem_id="custom_header")
+    gr.Markdown("CZ3004 is a module in Nanyang Technological University's Computer Science curriculum that involves creating a robot car that can navigate within an arena and around obstacles. Part of the assessment is to go to obstacles and detect alphanumeric symbols pasted on them.", elem_id="custom_title")
+    gr.Markdown("The two models available, Week 8 and Week 9, are for different subtasks. Week 8 model (as assessment was done in Week 8 of the school semester), \
+                is able to detect all symbols seen in the first three example images below. Week 9 model is limited to just the bullseye, left and right arrow symbols. \
+                Additionally, Week 9 model has been further trained on extreme edge cases where there is harsh sunlight behind the symbol/obstacle (seen in some of the examples).", elem_id="custom_title")
+    gr.Markdown("Heuristics used are based on the task as of AY22-23 Semester 2's edition of MDP. These include ignoring the bullseye symbol, taking only the biggest bounding box, and filtering similar sized detections by the expected position of the symbol based on where the robot is supposed to be relative to the symbol.", elem_id="custom_title")
+    gr.Markdown("This demo is part of a guide that is currently work-in-progress, for future CZ3004/SC2079 students to refer to.", elem_id="custom_title")
+
+
+    with gr.Row():
+            with gr.Column():
+                with gr.Box():
+                    gr.Markdown("## Inputs", elem_id="custom_header")
+                    input_image = gr.inputs.Image(type='pil', label="Original Image")
+                    btn = gr.Button(value="Submit")
+                    btn.style(full_width=True)
+            with gr.Column():
+                with gr.Box():
+                    gr.Markdown("## Parameters", elem_id="custom_header")
+                    model_selection = gr.inputs.Radio(['Week 8', 'Week 9'], type="value", default='Week 8', label='Model Selection')
+                    toggles = gr.CheckboxGroup(["Ignore Bullseye", "Biggest BBox Only and Position-Based Heuristics",], value=["Ignore Bullseye", "Biggest BBox Only and Position-Based Heuristics"], label="Heuristic Toggles")
+                    radios = gr.inputs.Radio(['Left', 'Center', 'Right', 'Disabled'], type="value", default='Center', label='Position Heuristic')
+    with gr.Row():
+            with gr.Box():
+                with gr.Column():
+                    gr.Markdown("## Raw Outputs", elem_id="custom_header")
+                    output_image = gr.outputs.Image(type="pil", label="Output Image")
+                    output_json = gr.outputs.JSON(label="Output JSON")
+            with gr.Box():
+                with gr.Column():
+                    gr.Markdown("## Filtered Outputs", elem_id="custom_header")
+                    filtered_image = gr.outputs.Image(type="pil", label="Filtered Output Image")
+                    filtered_json = gr.outputs.JSON(label="Filtered Output JSON")
+    with gr.Row():
+        gr.Examples(examples=examples,
+            inputs=input_image,
+            outputs=output_image,
+            fn=yolo,
+            cache_examples=False)
+    btn.click(yolo, inputs=[input_image, model_selection,toggles, radios], outputs=[output_image, output_json, filtered_image, filtered_json])
+
+demo.launch(debug=True)